huskies/server/src/chat/transport/matrix/bot/history.rs

//! Matrix conversation history — per-room message history for LLM context.
use crate::slog;
use matrix_sdk::ruma::OwnedRoomId;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::Mutex as TokioMutex;

/// Role of a participant in the conversation history.
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum ConversationRole {
    /// A message sent by a Matrix room participant.
    User,
    /// A response generated by the bot / LLM.
    Assistant,
}

/// A single turn in the per-room conversation history.
#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct ConversationEntry {
    pub role: ConversationRole,
    /// Matrix user ID (e.g. `@alice:example.com`).  Empty for assistant turns.
    pub sender: String,
    pub content: String,
}

/// Per-room state: conversation entries plus the Claude Code session ID for
/// structured conversation resumption.
#[derive(Clone, Debug, Default, Serialize, Deserialize)]
pub struct RoomConversation {
    /// Claude Code session ID used to resume multi-turn conversations so the
    /// LLM receives prior turns as structured API messages rather than a
    /// flattened text prefix.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub session_id: Option<String>,
    /// Rolling conversation entries (used for turn counting and persistence).
    pub entries: Vec<ConversationEntry>,
}

/// Per-room conversation state, keyed by room ID (serialised as string).
///
/// Wrapped in `Arc<TokioMutex<…>>` so it can be shared across concurrent
/// event-handler tasks without blocking the sync loop.
pub type ConversationHistory = Arc<TokioMutex<HashMap<OwnedRoomId, RoomConversation>>>;

/// On-disk format for persisted conversation history.  Room IDs are stored as
/// strings because `OwnedRoomId` does not implement `Serialize` as a map key.
#[derive(Serialize, Deserialize)]
pub(super) struct PersistedHistory {
    pub rooms: HashMap<String, RoomConversation>,
}

/// Path to the persisted conversation history file relative to project root.
pub(super) const HISTORY_FILE: &str = ".huskies/matrix_history.json";

/// Load conversation history from disk, returning an empty map on any error.
pub fn load_history(project_root: &std::path::Path) -> HashMap<OwnedRoomId, RoomConversation> {
    let path = project_root.join(HISTORY_FILE);
    let data = match std::fs::read_to_string(&path) {
        Ok(d) => d,
        Err(_) => return HashMap::new(),
    };
    let persisted: PersistedHistory = match serde_json::from_str(&data) {
        Ok(p) => p,
        Err(e) => {
            slog!("[matrix-bot] Failed to parse history file: {e}");
            return HashMap::new();
        }
    };
    persisted
        .rooms
        .into_iter()
        .filter_map(|(k, v)| k.parse::<OwnedRoomId>().ok().map(|room_id| (room_id, v)))
        .collect()
}

/// Save conversation history to disk.  Errors are logged but not propagated.
pub fn save_history(
    project_root: &std::path::Path,
    history: &HashMap<OwnedRoomId, RoomConversation>,
) {
    let persisted = PersistedHistory {
        rooms: history
            .iter()
            .map(|(k, v)| (k.to_string(), v.clone()))
            .collect(),
    };
    let path = project_root.join(HISTORY_FILE);
    match serde_json::to_string_pretty(&persisted) {
        Ok(json) => {
            if let Err(e) = std::fs::write(&path, json) {
                slog!("[matrix-bot] Failed to write history file: {e}");
            }
        }
        Err(e) => slog!("[matrix-bot] Failed to serialise history: {e}"),
    }
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

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

    #[tokio::test]
    async fn history_trims_to_configured_size() {
        let history: ConversationHistory = Arc::new(TokioMutex::new(HashMap::new()));
        let room_id: OwnedRoomId = "!test:example.com".parse().unwrap();
        let history_size = 4usize; // keep at most 4 entries

        // Add 6 entries (3 user + 3 assistant turns).
        {
            let mut guard = history.lock().await;
            let conv = guard.entry(room_id.clone()).or_default();
            conv.session_id = Some("test-session".to_string());
            for i in 0..3usize {
                conv.entries.push(ConversationEntry {
                    role: ConversationRole::User,
                    sender: "@user:example.com".to_string(),
                    content: format!("msg {i}"),
                });
                conv.entries.push(ConversationEntry {
                    role: ConversationRole::Assistant,
                    sender: String::new(),
                    content: format!("reply {i}"),
                });
                if conv.entries.len() > history_size {
                    let excess = conv.entries.len() - history_size;
                    conv.entries.drain(..excess);
                    conv.session_id = None;
                }
            }
        }

        let guard = history.lock().await;
        let conv = guard.get(&room_id).unwrap();
        assert_eq!(
            conv.entries.len(),
            history_size,
            "history must be trimmed to history_size"
        );
        // The oldest entries (msg 0 / reply 0) should have been dropped.
        assert!(
            conv.entries.iter().all(|e| !e.content.contains("msg 0")),
            "oldest entries must be dropped"
        );
        // Session ID must be cleared when trimming occurs.
        assert!(
            conv.session_id.is_none(),
            "session_id must be cleared on trim to start a fresh session"
        );
    }

    #[tokio::test]
    async fn each_room_has_independent_history() {
        let history: ConversationHistory = Arc::new(TokioMutex::new(HashMap::new()));
        let room_a: OwnedRoomId = "!room_a:example.com".parse().unwrap();
        let room_b: OwnedRoomId = "!room_b:example.com".parse().unwrap();

        {
            let mut guard = history.lock().await;
            guard
                .entry(room_a.clone())
                .or_default()
                .entries
                .push(ConversationEntry {
                    role: ConversationRole::User,
                    sender: "@alice:example.com".to_string(),
                    content: "Room A message".to_string(),
                });
            guard
                .entry(room_b.clone())
                .or_default()
                .entries
                .push(ConversationEntry {
                    role: ConversationRole::User,
                    sender: "@bob:example.com".to_string(),
                    content: "Room B message".to_string(),
                });
        }

        let guard = history.lock().await;
        let conv_a = guard.get(&room_a).unwrap();
        let conv_b = guard.get(&room_b).unwrap();
        assert_eq!(conv_a.entries.len(), 1);
        assert_eq!(conv_b.entries.len(), 1);
        assert_eq!(conv_a.entries[0].content, "Room A message");
        assert_eq!(conv_b.entries[0].content, "Room B message");
    }

    #[test]
    fn save_and_load_history_round_trip() {
        let dir = tempfile::tempdir().unwrap();
        let story_kit_dir = dir.path().join(".huskies");
        std::fs::create_dir_all(&story_kit_dir).unwrap();

        let room_id: OwnedRoomId = "!persist:example.com".parse().unwrap();
        let mut map: HashMap<OwnedRoomId, RoomConversation> = HashMap::new();
        let conv = map.entry(room_id.clone()).or_default();
        conv.session_id = Some("session-abc".to_string());
        conv.entries.push(ConversationEntry {
            role: ConversationRole::User,
            sender: "@alice:example.com".to_string(),
            content: "hello".to_string(),
        });
        conv.entries.push(ConversationEntry {
            role: ConversationRole::Assistant,
            sender: String::new(),
            content: "hi there!".to_string(),
        });

        save_history(dir.path(), &map);

        let loaded = load_history(dir.path());
        let loaded_conv = loaded.get(&room_id).expect("room must exist after load");
        assert_eq!(loaded_conv.session_id.as_deref(), Some("session-abc"));
        assert_eq!(loaded_conv.entries.len(), 2);
        assert_eq!(loaded_conv.entries[0].role, ConversationRole::User);
        assert_eq!(loaded_conv.entries[0].sender, "@alice:example.com");
        assert_eq!(loaded_conv.entries[0].content, "hello");
        assert_eq!(loaded_conv.entries[1].role, ConversationRole::Assistant);
        assert_eq!(loaded_conv.entries[1].content, "hi there!");
    }

    #[test]
    fn load_history_returns_empty_on_missing_file() {
        let dir = tempfile::tempdir().unwrap();
        let loaded = load_history(dir.path());
        assert!(loaded.is_empty());
    }

    #[test]
    fn load_history_returns_empty_on_corrupt_file() {
        let dir = tempfile::tempdir().unwrap();
        let story_kit_dir = dir.path().join(".huskies");
        std::fs::create_dir_all(&story_kit_dir).unwrap();
        std::fs::write(dir.path().join(HISTORY_FILE), "not valid json").unwrap();
        let loaded = load_history(dir.path());
        assert!(loaded.is_empty());
    }

    #[tokio::test]
    async fn session_id_preserved_within_history_size() {
        let history: ConversationHistory = Arc::new(TokioMutex::new(HashMap::new()));
        let room_id: OwnedRoomId = "!session:example.com".parse().unwrap();

        {
            let mut guard = history.lock().await;
            let conv = guard.entry(room_id.clone()).or_default();
            conv.session_id = Some("sess-1".to_string());
            conv.entries.push(ConversationEntry {
                role: ConversationRole::User,
                sender: "@alice:example.com".to_string(),
                content: "hello".to_string(),
            });
            conv.entries.push(ConversationEntry {
                role: ConversationRole::Assistant,
                sender: String::new(),
                content: "hi".to_string(),
            });
            // No trimming needed (2 entries, well under any reasonable limit).
        }

        let guard = history.lock().await;
        let conv = guard.get(&room_id).unwrap();
        assert_eq!(
            conv.session_id.as_deref(),
            Some("sess-1"),
            "session_id must be preserved when no trimming occurs"
        );
    }

    #[tokio::test]
    async fn multi_user_entries_preserve_sender() {
        let history: ConversationHistory = Arc::new(TokioMutex::new(HashMap::new()));
        let room_id: OwnedRoomId = "!multi:example.com".parse().unwrap();

        {
            let mut guard = history.lock().await;
            let conv = guard.entry(room_id.clone()).or_default();
            conv.entries.push(ConversationEntry {
                role: ConversationRole::User,
                sender: "@alice:example.com".to_string(),
                content: "from alice".to_string(),
            });
            conv.entries.push(ConversationEntry {
                role: ConversationRole::Assistant,
                sender: String::new(),
                content: "reply to alice".to_string(),
            });
            conv.entries.push(ConversationEntry {
                role: ConversationRole::User,
                sender: "@bob:example.com".to_string(),
                content: "from bob".to_string(),
            });
        }

        let guard = history.lock().await;
        let conv = guard.get(&room_id).unwrap();
        assert_eq!(conv.entries[0].sender, "@alice:example.com");
        assert_eq!(conv.entries[2].sender, "@bob:example.com");
    }
}