use chrono::Utc;
use entelix_core::ir::{ContentPart, Message, Role};
use serde::{Deserialize, Serialize};
use crate::event::GraphEvent;
#[derive(Clone, Debug, Serialize, Deserialize)]
#[non_exhaustive]
pub struct SessionGraph {
pub thread_id: String,
pub events: Vec<GraphEvent>,
archival_watermark: usize,
}
impl SessionGraph {
pub fn new(thread_id: impl Into<String>) -> Self {
Self {
thread_id: thread_id.into(),
events: Vec::new(),
archival_watermark: 0,
}
}
pub fn append(&mut self, event: GraphEvent) -> usize {
self.events.push(event);
self.events.len().saturating_sub(1)
}
pub const fn len(&self) -> usize {
self.events.len()
}
pub const fn is_empty(&self) -> bool {
self.events.is_empty()
}
pub fn events_since(&self, cursor: usize) -> &[GraphEvent] {
self.events.get(cursor..).unwrap_or(&[])
}
pub fn replay_into<S, R>(&self, initial: S, mut reducer: R) -> S
where
R: FnMut(&mut S, &GraphEvent),
{
let mut state = initial;
for event in &self.events {
reducer(&mut state, event);
}
state
}
pub fn current_branch_messages(&self) -> Vec<Message> {
let mut out = Vec::new();
for event in &self.events {
match event {
GraphEvent::UserMessage { content, .. } => {
out.push(Message::new(Role::User, content.clone()));
}
GraphEvent::AssistantMessage { content, .. } => {
out.push(Message::new(Role::Assistant, content.clone()));
}
GraphEvent::ToolResult {
tool_use_id,
name,
content,
is_error,
..
} => out.push(Message::new(
Role::Tool,
vec![ContentPart::ToolResult {
tool_use_id: tool_use_id.clone(),
name: name.clone(),
content: content.clone(),
is_error: *is_error,
cache_control: None,
provider_echoes: Vec::new(),
}],
)),
_ => {}
}
}
out
}
pub fn fork(&mut self, branch_at: usize, new_thread_id: impl Into<String>) -> Option<Self> {
let cloned_events = match self.events.get(..=branch_at) {
Some(slice) => slice.to_vec(),
None => return None,
};
let new_thread_id = new_thread_id.into();
self.append(GraphEvent::BranchCreated {
branch_id: new_thread_id.clone(),
parent_event: branch_at,
timestamp: Utc::now(),
});
Some(Self {
thread_id: new_thread_id,
events: cloned_events,
archival_watermark: 0,
})
}
pub const fn archive_before(&mut self, watermark: usize) {
if watermark > self.archival_watermark && watermark <= self.events.len() {
self.archival_watermark = watermark;
}
}
pub const fn archival_watermark(&self) -> usize {
self.archival_watermark
}
pub fn branch_events(&self) -> impl Iterator<Item = (&str, usize)> {
self.events.iter().filter_map(|e| match e {
GraphEvent::BranchCreated {
branch_id,
parent_event,
..
} => Some((branch_id.as_str(), *parent_event)),
_ => None,
})
}
}