#![cfg(feature = "llm")]
use std::collections::BTreeMap;
use std::path::PathBuf;
use std::sync::Mutex;
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use tokio::io::{AsyncBufReadExt, AsyncWriteExt};
use crate::error::AgentError;
use crate::llm::Message;
fn default_version() -> u32 {
1
}
fn new_uuid() -> String {
uuid::Uuid::new_v4().to_string()
}
fn now_rfc3339() -> String {
chrono::Utc::now().to_rfc3339()
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct SessionMeta {
pub id: String,
#[serde(default = "default_version")]
pub version: u32,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub parent_session: Option<String>,
#[serde(default)]
pub created_at: String,
}
impl SessionMeta {
pub fn new() -> Self {
Self {
id: new_uuid(),
version: 1,
parent_session: None,
created_at: now_rfc3339(),
}
}
}
impl Default for SessionMeta {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum SessionEntry {
Message {
id: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
parent_id: Option<String>,
timestamp: String,
message: Message,
},
Compaction {
id: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
parent_id: Option<String>,
timestamp: String,
summary: String,
first_kept_id: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
tokens_before: Option<u64>,
},
}
impl SessionEntry {
pub fn message(message: Message) -> Self {
SessionEntry::Message {
id: new_uuid(),
parent_id: None,
timestamp: now_rfc3339(),
message,
}
}
pub fn compaction(summary: String, first_kept_id: String, tokens_before: Option<u64>) -> Self {
SessionEntry::Compaction {
id: new_uuid(),
parent_id: None,
timestamp: now_rfc3339(),
summary,
first_kept_id,
tokens_before,
}
}
pub fn id(&self) -> &str {
match self {
SessionEntry::Message { id, .. } => id,
SessionEntry::Compaction { id, .. } => id,
}
}
}
#[async_trait]
pub trait SessionStore: Send + Sync {
async fn create(&self, meta: SessionMeta) -> Result<String, AgentError>;
async fn append(&self, session_id: &str, entry: SessionEntry) -> Result<(), AgentError>;
async fn load(&self, session_id: &str) -> Result<Vec<SessionEntry>, AgentError>;
async fn list(&self) -> Result<Vec<SessionMeta>, AgentError>;
}
fn session_not_found(session_id: &str) -> AgentError {
AgentError::Other(format!("Session not found: {session_id}"))
}
type SessionMap = BTreeMap<String, (SessionMeta, Vec<SessionEntry>)>;
#[derive(Debug, Default)]
pub struct InMemorySessionStore {
sessions: Mutex<SessionMap>,
}
impl InMemorySessionStore {
pub fn new() -> Self {
Self::default()
}
fn lock(&self) -> std::sync::MutexGuard<'_, SessionMap> {
self.sessions
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner())
}
}
#[async_trait]
impl SessionStore for InMemorySessionStore {
async fn create(&self, meta: SessionMeta) -> Result<String, AgentError> {
let mut sessions = self.lock();
let id = meta.id.clone();
if sessions.contains_key(&id) {
return Err(AgentError::DuplicateId(id));
}
sessions.insert(id.clone(), (meta, Vec::new()));
Ok(id)
}
async fn append(&self, session_id: &str, entry: SessionEntry) -> Result<(), AgentError> {
let mut sessions = self.lock();
let (_, entries) = sessions
.get_mut(session_id)
.ok_or_else(|| session_not_found(session_id))?;
entries.push(entry);
Ok(())
}
async fn load(&self, session_id: &str) -> Result<Vec<SessionEntry>, AgentError> {
let sessions = self.lock();
let (_, entries) = sessions
.get(session_id)
.ok_or_else(|| session_not_found(session_id))?;
Ok(entries.clone())
}
async fn list(&self) -> Result<Vec<SessionMeta>, AgentError> {
let sessions = self.lock();
Ok(sessions.values().map(|(meta, _)| meta.clone()).collect())
}
}
#[derive(Debug, Clone)]
pub struct JsonlSessionStore {
dir: PathBuf,
}
impl JsonlSessionStore {
pub fn new(dir: impl Into<PathBuf>) -> Self {
Self { dir: dir.into() }
}
fn session_path(&self, session_id: &str) -> PathBuf {
self.dir.join(format!("{session_id}.jsonl"))
}
}
#[async_trait]
impl SessionStore for JsonlSessionStore {
async fn create(&self, meta: SessionMeta) -> Result<String, AgentError> {
tokio::fs::create_dir_all(&self.dir)
.await
.map_err(|e| AgentError::IoError(format!("Failed to create session dir: {e}")))?;
let id = meta.id.clone();
let path = self.session_path(&id);
let mut file = tokio::fs::OpenOptions::new()
.write(true)
.create_new(true)
.open(&path)
.await
.map_err(|e| {
if e.kind() == std::io::ErrorKind::AlreadyExists {
AgentError::DuplicateId(id.clone())
} else {
AgentError::IoError(format!("Failed to create session file: {e}"))
}
})?;
let mut line = serde_json::to_string(&meta).map_err(|e| {
AgentError::SerializationError(format!("Failed to serialize meta: {e}"))
})?;
line.push('\n');
file.write_all(line.as_bytes())
.await
.map_err(|e| AgentError::IoError(format!("Failed to write session header: {e}")))?;
file.flush()
.await
.map_err(|e| AgentError::IoError(format!("Failed to flush session file: {e}")))?;
Ok(id)
}
async fn append(&self, session_id: &str, entry: SessionEntry) -> Result<(), AgentError> {
let path = self.session_path(session_id);
let mut file = tokio::fs::OpenOptions::new()
.append(true)
.open(&path)
.await
.map_err(|e| {
if e.kind() == std::io::ErrorKind::NotFound {
session_not_found(session_id)
} else {
AgentError::IoError(format!("Failed to open session file: {e}"))
}
})?;
let mut line = serde_json::to_string(&entry).map_err(|e| {
AgentError::SerializationError(format!("Failed to serialize entry: {e}"))
})?;
line.push('\n');
file.write_all(line.as_bytes())
.await
.map_err(|e| AgentError::IoError(format!("Failed to append session entry: {e}")))?;
file.flush()
.await
.map_err(|e| AgentError::IoError(format!("Failed to flush session file: {e}")))?;
Ok(())
}
async fn load(&self, session_id: &str) -> Result<Vec<SessionEntry>, AgentError> {
let path = self.session_path(session_id);
let content = tokio::fs::read_to_string(&path).await.map_err(|e| {
if e.kind() == std::io::ErrorKind::NotFound {
session_not_found(session_id)
} else {
AgentError::IoError(format!("Failed to read session file: {e}"))
}
})?;
let complete_len = if content.ends_with('\n') {
content.len()
} else {
content.rfind('\n').map(|i| i + 1).unwrap_or(0)
};
if complete_len < content.len() {
log::warn!("Discarding crash-truncated final line in session {session_id}");
let file = tokio::fs::OpenOptions::new()
.write(true)
.open(&path)
.await
.map_err(|e| {
AgentError::IoError(format!("Failed to open session file for repair: {e}"))
})?;
file.set_len(complete_len as u64).await.map_err(|e| {
AgentError::IoError(format!("Failed to truncate session file: {e}"))
})?;
}
let mut lines = content[..complete_len].lines();
let Some(header) = lines.next() else {
return Err(AgentError::JsonParseError(format!(
"Session file for {session_id} is empty (missing header line)"
)));
};
serde_json::from_str::<SessionMeta>(header).map_err(|e| {
AgentError::JsonParseError(format!("Invalid session header for {session_id}: {e}"))
})?;
let mut entries = Vec::new();
for (i, line) in lines.enumerate() {
let entry = serde_json::from_str::<SessionEntry>(line).map_err(|e| {
AgentError::JsonParseError(format!(
"Invalid entry at line {} of session {session_id}: {e}",
i + 2
))
})?;
entries.push(entry);
}
Ok(entries)
}
async fn list(&self) -> Result<Vec<SessionMeta>, AgentError> {
let mut read_dir = match tokio::fs::read_dir(&self.dir).await {
Ok(rd) => rd,
Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(Vec::new()),
Err(e) => {
return Err(AgentError::IoError(format!(
"Failed to read session dir: {e}"
)));
}
};
let mut metas = Vec::new();
loop {
let entry = match read_dir.next_entry().await {
Ok(Some(entry)) => entry,
Ok(None) => break,
Err(e) => {
return Err(AgentError::IoError(format!(
"Failed to read session dir entry: {e}"
)));
}
};
let path = entry.path();
if path.extension().and_then(|e| e.to_str()) != Some("jsonl") {
continue;
}
let file = match tokio::fs::File::open(&path).await {
Ok(file) => file,
Err(e) => {
log::warn!("Skipping unreadable session file {}: {e}", path.display());
continue;
}
};
let mut header = String::new();
match tokio::io::BufReader::new(file).read_line(&mut header).await {
Ok(0) => {
log::warn!("Skipping empty session file {}", path.display());
continue;
}
Ok(_) => {}
Err(e) => {
log::warn!("Skipping unreadable session file {}: {e}", path.display());
continue;
}
}
match serde_json::from_str::<SessionMeta>(header.trim_end()) {
Ok(meta) => metas.push(meta),
Err(e) => {
log::warn!(
"Skipping session file with invalid header {}: {e}",
path.display()
);
}
}
}
Ok(metas)
}
}
pub fn build_context(entries: &[SessionEntry]) -> Vec<Message> {
let last_compaction = entries.iter().enumerate().rev().find_map(|(i, e)| match e {
SessionEntry::Compaction {
summary,
first_kept_id,
..
} => Some((i, summary, first_kept_id)),
_ => None,
});
let Some((compaction_index, summary, first_kept_id)) = last_compaction else {
return entries.iter().filter_map(entry_message).collect();
};
let mut context = vec![Message::user(format!("[Conversation summary]\n{summary}"))];
let start = entries
.iter()
.position(|e| matches!(e, SessionEntry::Message { id, .. } if id == first_kept_id))
.unwrap_or(compaction_index + 1);
context.extend(entries[start..].iter().filter_map(entry_message));
context
}
fn entry_message(entry: &SessionEntry) -> Option<Message> {
match entry {
SessionEntry::Message { message, .. } => Some(message.clone()),
_ => None,
}
}
#[cfg(test)]
mod tests {
use super::*;
fn message_entry(id: &str, text: &str) -> SessionEntry {
SessionEntry::Message {
id: id.to_string(),
parent_id: None,
timestamp: now_rfc3339(),
message: Message::user(text.to_string()),
}
}
fn compaction_entry(id: &str, summary: &str, first_kept_id: &str) -> SessionEntry {
SessionEntry::Compaction {
id: id.to_string(),
parent_id: None,
timestamp: now_rfc3339(),
summary: summary.to_string(),
first_kept_id: first_kept_id.to_string(),
tokens_before: Some(1000),
}
}
#[tokio::test]
async fn test_in_memory_round_trip() {
let store = InMemorySessionStore::new();
let meta = SessionMeta::new();
let id = store.create(meta.clone()).await.unwrap();
assert_eq!(id, meta.id);
let e1 = message_entry("e1", "hello");
let e2 = message_entry("e2", "world");
store.append(&id, e1.clone()).await.unwrap();
store.append(&id, e2.clone()).await.unwrap();
let entries = store.load(&id).await.unwrap();
assert_eq!(entries, vec![e1, e2]);
let metas = store.list().await.unwrap();
assert_eq!(metas, vec![meta]);
}
#[tokio::test]
async fn test_in_memory_duplicate_create_errors() {
let store = InMemorySessionStore::new();
let meta = SessionMeta::new();
store.create(meta.clone()).await.unwrap();
assert!(store.create(meta).await.is_err());
}
#[tokio::test]
async fn test_in_memory_unknown_session_errors() {
let store = InMemorySessionStore::new();
assert!(store.load("missing").await.is_err());
assert!(
store
.append("missing", message_entry("e1", "x"))
.await
.is_err()
);
}
#[tokio::test]
async fn test_jsonl_replay_and_list_with_new_store() {
let dir = tempfile::tempdir().unwrap();
let store = JsonlSessionStore::new(dir.path());
let meta = SessionMeta::new();
let id = store.create(meta.clone()).await.unwrap();
let e1 = message_entry("e1", "hello");
let e2 = SessionEntry::compaction("summary".to_string(), "e1".to_string(), Some(42));
store.append(&id, e1.clone()).await.unwrap();
store.append(&id, e2.clone()).await.unwrap();
let reopened = JsonlSessionStore::new(dir.path());
let entries = reopened.load(&id).await.unwrap();
assert_eq!(entries, vec![e1, e2]);
let metas = reopened.list().await.unwrap();
assert_eq!(metas, vec![meta]);
}
#[tokio::test]
async fn test_jsonl_duplicate_create_errors() {
let dir = tempfile::tempdir().unwrap();
let store = JsonlSessionStore::new(dir.path());
let meta = SessionMeta::new();
store.create(meta.clone()).await.unwrap();
assert!(store.create(meta).await.is_err());
}
#[tokio::test]
async fn test_jsonl_truncated_final_line_is_discarded_and_repaired() {
let dir = tempfile::tempdir().unwrap();
let store = JsonlSessionStore::new(dir.path());
let meta = SessionMeta::new();
let id = store.create(meta).await.unwrap();
let e1 = message_entry("e1", "hello");
store.append(&id, e1.clone()).await.unwrap();
let path = dir.path().join(format!("{id}.jsonl"));
let mut content = std::fs::read_to_string(&path).unwrap();
let intact_len = content.len();
content.push_str("{\"type\":\"message\",\"id\":\"e2");
std::fs::write(&path, content).unwrap();
let entries = store.load(&id).await.unwrap();
assert_eq!(entries, vec![e1.clone()]);
let repaired = std::fs::read_to_string(&path).unwrap();
assert_eq!(repaired.len(), intact_len);
assert!(repaired.ends_with('\n'));
let e2 = message_entry("e2", "world");
store.append(&id, e2.clone()).await.unwrap();
let e3 = message_entry("e3", "again");
store.append(&id, e3.clone()).await.unwrap();
let entries = store.load(&id).await.unwrap();
assert_eq!(entries, vec![e1, e2, e3]);
}
#[tokio::test]
async fn test_jsonl_unparseable_complete_final_line_errors() {
let dir = tempfile::tempdir().unwrap();
let store = JsonlSessionStore::new(dir.path());
let meta = SessionMeta::new();
let id = store.create(meta).await.unwrap();
store
.append(&id, message_entry("e1", "hello"))
.await
.unwrap();
let path = dir.path().join(format!("{id}.jsonl"));
let mut content = std::fs::read_to_string(&path).unwrap();
content.push_str("not json\n");
std::fs::write(&path, content).unwrap();
assert!(store.load(&id).await.is_err());
}
#[tokio::test]
async fn test_jsonl_unparseable_middle_line_errors() {
let dir = tempfile::tempdir().unwrap();
let store = JsonlSessionStore::new(dir.path());
let meta = SessionMeta::new();
let id = store.create(meta).await.unwrap();
store
.append(&id, message_entry("e1", "hello"))
.await
.unwrap();
let path = dir.path().join(format!("{id}.jsonl"));
let content = std::fs::read_to_string(&path).unwrap();
let lines: Vec<&str> = content.lines().collect();
let corrupted = format!("{}\nnot json\n{}\n", lines[0], lines[1]);
std::fs::write(&path, corrupted).unwrap();
assert!(store.load(&id).await.is_err());
}
#[tokio::test]
async fn test_jsonl_unknown_session_errors_and_empty_dir_lists_nothing() {
let dir = tempfile::tempdir().unwrap();
let store = JsonlSessionStore::new(dir.path().join("never-created"));
assert!(store.load("missing").await.is_err());
assert!(
store
.append("missing", message_entry("e1", "x"))
.await
.is_err()
);
assert_eq!(store.list().await.unwrap(), vec![]);
}
#[test]
fn test_build_context_no_compaction() {
let entries = vec![message_entry("e1", "a"), message_entry("e2", "b")];
let context = build_context(&entries);
assert_eq!(context.len(), 2);
assert_eq!(context[0].text(), "a");
assert_eq!(context[1].text(), "b");
}
#[test]
fn test_build_context_with_compaction_drops_prefix() {
let entries = vec![
message_entry("e1", "old1"),
message_entry("e2", "old2"),
message_entry("e3", "kept"),
message_entry("e4", "tail"),
compaction_entry("c1", "the summary", "e3"),
];
let context = build_context(&entries);
assert_eq!(context.len(), 3);
assert_eq!(context[0].role, "user");
assert_eq!(context[0].text(), "[Conversation summary]\nthe summary");
assert_eq!(context[1].text(), "kept");
assert_eq!(context[2].text(), "tail");
}
#[test]
fn test_build_context_last_compaction_wins() {
let entries = vec![
message_entry("e1", "a"),
compaction_entry("c1", "first summary", "e1"),
message_entry("e2", "b"),
message_entry("e3", "c"),
compaction_entry("c2", "second summary", "e3"),
message_entry("e4", "d"),
];
let context = build_context(&entries);
assert_eq!(context.len(), 3);
assert_eq!(context[0].text(), "[Conversation summary]\nsecond summary");
assert_eq!(context[1].text(), "c");
assert_eq!(context[2].text(), "d");
}
#[test]
fn test_build_context_first_kept_id_not_found_falls_back() {
let entries = vec![
message_entry("e1", "old"),
compaction_entry("c1", "the summary", "nonexistent"),
message_entry("e2", "after1"),
message_entry("e3", "after2"),
];
let context = build_context(&entries);
assert_eq!(context.len(), 3);
assert_eq!(context[0].text(), "[Conversation summary]\nthe summary");
assert_eq!(context[1].text(), "after1");
assert_eq!(context[2].text(), "after2");
}
#[test]
fn test_session_entry_message_serde_round_trip() {
let entry = message_entry("e1", "hello");
let json = serde_json::to_value(&entry).unwrap();
assert_eq!(json["type"], serde_json::json!("message"));
assert!(json.get("parent_id").is_none());
let restored: SessionEntry = serde_json::from_value(json).unwrap();
assert_eq!(restored, entry);
}
#[test]
fn test_session_entry_compaction_serde_round_trip() {
let entry = compaction_entry("c1", "summary", "e1");
let json = serde_json::to_value(&entry).unwrap();
assert_eq!(json["type"], serde_json::json!("compaction"));
assert_eq!(json["tokens_before"], serde_json::json!(1000));
let restored: SessionEntry = serde_json::from_value(json).unwrap();
assert_eq!(restored, entry);
}
#[test]
fn test_session_entry_unknown_field_ignored() {
let json = serde_json::json!({
"type": "message",
"id": "e1",
"timestamp": "2026-07-19T00:00:00+00:00",
"message": { "role": "user", "content": "hi" },
"some_future_field": true,
});
let entry: SessionEntry = serde_json::from_value(json).unwrap();
assert_eq!(entry.id(), "e1");
}
#[test]
fn test_session_meta_serde_round_trip_and_forward_compat() {
let meta = SessionMeta::new();
let json = serde_json::to_value(&meta).unwrap();
assert!(json.get("parent_session").is_none());
let restored: SessionMeta = serde_json::from_value(json).unwrap();
assert_eq!(restored, meta);
let json = serde_json::json!({
"id": "s1",
"some_future_field": "x",
});
let meta: SessionMeta = serde_json::from_value(json).unwrap();
assert_eq!(meta.id, "s1");
assert_eq!(meta.version, 1);
assert_eq!(meta.parent_session, None);
assert_eq!(meta.created_at, "");
}
#[test]
fn test_generated_ids_are_unique() {
let m1 = SessionMeta::new();
let m2 = SessionMeta::new();
assert_ne!(m1.id, m2.id);
let e1 = SessionEntry::message(Message::user("a".to_string()));
let e2 = SessionEntry::message(Message::user("a".to_string()));
assert_ne!(e1.id(), e2.id());
}
}