use std::collections::HashMap;
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Session {
pub key: String,
#[serde(default)]
pub messages: Vec<serde_json::Value>,
#[serde(default = "Utc::now")]
pub created_at: DateTime<Utc>,
#[serde(default = "Utc::now")]
pub updated_at: DateTime<Utc>,
#[serde(default)]
pub metadata: HashMap<String, serde_json::Value>,
#[serde(default)]
pub last_consolidated: usize,
}
impl Session {
pub fn new(key: impl Into<String>) -> Self {
let now = Utc::now();
Self {
key: key.into(),
messages: Vec::new(),
created_at: now,
updated_at: now,
metadata: HashMap::new(),
last_consolidated: 0,
}
}
pub fn add_message(
&mut self,
role: &str,
content: &str,
extras: Option<HashMap<String, serde_json::Value>>,
) {
let mut msg = serde_json::json!({
"role": role,
"content": content,
"timestamp": Utc::now().to_rfc3339(),
});
if let Some(extras) = extras
&& let Some(obj) = msg.as_object_mut()
{
for (k, v) in extras {
obj.insert(k, v);
}
}
self.messages.push(msg);
self.updated_at = Utc::now();
}
pub fn get_history(&self, max_messages: usize) -> Vec<serde_json::Value> {
let start = self.messages.len().saturating_sub(max_messages);
self.messages[start..]
.iter()
.map(|m| {
let mut msg = serde_json::json!({
"role": m.get("role").and_then(|v| v.as_str()).unwrap_or("user"),
"content": m.get("content").and_then(|v| v.as_str()).unwrap_or(""),
});
if let Some(tool_call_id) = m.get("tool_call_id").filter(|v| !v.is_null()) {
msg["tool_call_id"] = tool_call_id.clone();
}
if let Some(tool_calls) = m.get("tool_calls").filter(|v| !v.is_null()) {
msg["tool_calls"] = tool_calls.clone();
}
msg
})
.collect()
}
pub fn clear(&mut self) {
self.messages.clear();
self.last_consolidated = 0;
self.updated_at = Utc::now();
}
}
impl Default for Session {
fn default() -> Self {
Self::new("")
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn new_session() {
let s = Session::new("telegram:123");
assert_eq!(s.key, "telegram:123");
assert!(s.messages.is_empty());
assert_eq!(s.last_consolidated, 0);
}
#[test]
fn add_message_basic() {
let mut s = Session::new("test");
s.add_message("user", "hello", None);
s.add_message("assistant", "hi there", None);
assert_eq!(s.messages.len(), 2);
assert_eq!(s.messages[0]["role"], "user");
assert_eq!(s.messages[1]["content"], "hi there");
}
#[test]
fn add_message_with_extras() {
let mut s = Session::new("test");
let mut extras = HashMap::new();
extras.insert("tool_calls".into(), serde_json::json!([{"id": "tc1"}]));
s.add_message("assistant", "let me check", Some(extras));
assert!(s.messages[0].get("tool_calls").is_some());
}
#[test]
fn get_history_all() {
let mut s = Session::new("test");
s.add_message("user", "one", None);
s.add_message("assistant", "two", None);
let hist = s.get_history(500);
assert_eq!(hist.len(), 2);
assert_eq!(hist[0]["role"], "user");
assert_eq!(hist[1]["content"], "two");
}
#[test]
fn get_history_truncated() {
let mut s = Session::new("test");
for i in 0..10 {
s.add_message("user", &format!("msg {i}"), None);
}
let hist = s.get_history(3);
assert_eq!(hist.len(), 3);
assert_eq!(hist[0]["content"], "msg 7");
assert_eq!(hist[2]["content"], "msg 9");
}
#[test]
fn clear_resets_state() {
let mut s = Session::new("test");
s.add_message("user", "hello", None);
s.last_consolidated = 1;
s.clear();
assert!(s.messages.is_empty());
assert_eq!(s.last_consolidated, 0);
}
#[test]
fn serde_roundtrip() {
let mut s = Session::new("slack:C123");
s.add_message("user", "test", None);
s.last_consolidated = 0;
let json = serde_json::to_string(&s).unwrap();
let restored: Session = serde_json::from_str(&json).unwrap();
assert_eq!(restored.key, "slack:C123");
assert_eq!(restored.messages.len(), 1);
}
#[test]
fn default_session() {
let s = Session::default();
assert_eq!(s.key, "");
assert!(s.messages.is_empty());
}
}