310 lines
12 KiB
Rust
310 lines
12 KiB
Rust
|
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 = ".storkit/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(".storkit");
|
||
|
|
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(".storkit");
|
||
|
|
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");
|
||
|
|
}
|
||
|
|
}
|