use std::collections::{BTreeMap, BTreeSet};
use serde::{Deserialize, Serialize};
use crate::sidecar::NativeTurn;
use crate::{ChatMessage, InterchangeError as Error, Result};
pub type NodeId = String;
#[derive(Debug, Clone)]
pub struct TreeNode {
pub id: NodeId,
pub parent: Option<NodeId>,
pub children: Vec<NodeId>,
pub message: ChatMessage,
pub label: Option<String>,
pub created_at_ms: i64,
}
#[derive(Serialize, Deserialize)]
struct TreeNodeWire {
id: NodeId,
parent: Option<NodeId>,
#[serde(default)]
children: Vec<NodeId>,
message: NativeTurn,
#[serde(default)]
label: Option<String>,
#[serde(default)]
created_at_ms: i64,
}
impl From<&TreeNode> for TreeNodeWire {
fn from(n: &TreeNode) -> Self {
let message = NativeTurn {
supercode_turn: 1,
ts: crate::sidecar::ms_to_rfc3339(n.created_at_ms),
role: n.message.role,
content: n.message.content.clone(),
content_parts: n.message.content_parts.clone(),
tool_calls: n.message.tool_calls.clone(),
tool_call_id: n.message.tool_call_id.clone(),
name: n.message.name.clone(),
metadata: n.message.metadata.clone(),
};
TreeNodeWire {
id: n.id.clone(),
parent: n.parent.clone(),
children: n.children.clone(),
message,
label: n.label.clone(),
created_at_ms: n.created_at_ms,
}
}
}
impl From<TreeNodeWire> for TreeNode {
fn from(w: TreeNodeWire) -> Self {
TreeNode {
id: w.id,
parent: w.parent,
children: w.children,
message: w.message.into_message(),
label: w.label,
created_at_ms: w.created_at_ms,
}
}
}
impl Serialize for TreeNode {
fn serialize<S: serde::Serializer>(&self, ser: S) -> std::result::Result<S::Ok, S::Error> {
TreeNodeWire::from(self).serialize(ser)
}
}
impl<'de> Deserialize<'de> for TreeNode {
fn deserialize<D: serde::Deserializer<'de>>(de: D) -> std::result::Result<Self, D::Error> {
TreeNodeWire::deserialize(de).map(TreeNode::from)
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct BranchSummary {
pub summary: String,
pub node_id: NodeId,
pub branch: String,
#[serde(default)]
pub model_id: Option<String>,
#[serde(default)]
pub created_at_ms: i64,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Branch {
pub name: String,
pub leaf: Option<NodeId>,
#[serde(default)]
pub summary: Option<BranchSummary>,
#[serde(default)]
pub created_at_ms: i64,
}
pub const MAIN_BRANCH: &str = "main";
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SessionTree {
pub nodes: BTreeMap<NodeId, TreeNode>,
pub root: Option<NodeId>,
pub branches: BTreeMap<String, Branch>,
pub active_branch: String,
#[serde(default)]
next_id: u64,
}
impl Default for SessionTree {
fn default() -> Self {
Self::new()
}
}
impl SessionTree {
pub fn new() -> Self {
let mut branches = BTreeMap::new();
branches.insert(
MAIN_BRANCH.to_string(),
Branch {
name: MAIN_BRANCH.to_string(),
leaf: None,
summary: None,
created_at_ms: 0,
},
);
SessionTree {
nodes: BTreeMap::new(),
root: None,
branches,
active_branch: MAIN_BRANCH.to_string(),
next_id: 0,
}
}
pub fn from_linear(messages: &[ChatMessage], created_at_ms: i64) -> Self {
let mut tree = Self::new();
for m in messages {
tree.append_message(m.clone(), created_at_ms);
}
tree
}
fn alloc_id(&mut self) -> NodeId {
loop {
let id = format!("n{}", self.next_id);
self.next_id += 1;
if !self.nodes.contains_key(&id) {
return id;
}
}
}
pub fn node(&self, id: &str) -> Option<&TreeNode> {
self.nodes.get(id)
}
fn require_node(&self, id: &str) -> Result<&TreeNode> {
self.nodes
.get(id)
.ok_or_else(|| Error::Other(format!("session tree has no node `{id}`")))
}
fn require_branch(&self, name: &str) -> Result<&Branch> {
self.branches
.get(name)
.ok_or_else(|| Error::Other(format!("session tree has no branch `{name}`")))
}
pub fn append_message(&mut self, message: ChatMessage, created_at_ms: i64) -> NodeId {
let parent = self
.branches
.get(&self.active_branch)
.and_then(|b| b.leaf.clone());
let id = self.alloc_id();
self.nodes.insert(
id.clone(),
TreeNode {
id: id.clone(),
parent: parent.clone(),
children: Vec::new(),
message,
label: None,
created_at_ms,
},
);
match &parent {
Some(p) => {
if let Some(pn) = self.nodes.get_mut(p) {
pn.children.push(id.clone());
}
}
None => self.root = Some(id.clone()),
}
if let Some(b) = self.branches.get_mut(&self.active_branch) {
b.leaf = Some(id.clone());
}
id
}
fn fresh_branch_name(&self, base: &str) -> String {
if !self.branches.contains_key(base) {
return base.to_string();
}
let mut n = 2u64;
loop {
let candidate = format!("{base}-{n}");
if !self.branches.contains_key(&candidate) {
return candidate;
}
n += 1;
}
}
pub fn rewind(&mut self, node_id: &str, timestamp_ms: i64) -> Result<Option<String>> {
self.require_node(node_id)?;
let old_leaf = self
.branches
.get(&self.active_branch)
.and_then(|b| b.leaf.clone());
let preserved = match &old_leaf {
Some(old) if old != node_id => {
let name = self.fresh_branch_name(&format!("{}-rewound", self.active_branch));
self.branches.insert(
name.clone(),
Branch {
name: name.clone(),
leaf: Some(old.clone()),
summary: None,
created_at_ms: timestamp_ms,
},
);
Some(name)
}
_ => None,
};
if let Some(b) = self.branches.get_mut(&self.active_branch) {
b.leaf = Some(node_id.to_string());
}
Ok(preserved)
}
pub fn branch(
&mut self,
from_node: &str,
name: Option<String>,
timestamp_ms: i64,
) -> Result<String> {
self.require_node(from_node)?;
let name = match name {
Some(n) => {
if self.branches.contains_key(&n) {
return Err(Error::Other(format!(
"session tree already has a branch named `{n}`"
)));
}
n
}
None => self.fresh_branch_name("branch"),
};
self.branches.insert(
name.clone(),
Branch {
name: name.clone(),
leaf: Some(from_node.to_string()),
summary: None,
created_at_ms: timestamp_ms,
},
);
self.active_branch = name.clone();
Ok(name)
}
pub fn switch_branch(&mut self, name: &str) -> Result<()> {
self.require_branch(name)?;
self.active_branch = name.to_string();
Ok(())
}
pub fn label(&mut self, node_id: &str, label: impl Into<String>) -> Result<()> {
let node = self
.nodes
.get_mut(node_id)
.ok_or_else(|| Error::Other(format!("session tree has no node `{node_id}`")))?;
node.label = Some(label.into());
Ok(())
}
pub fn clear_label(&mut self, node_id: &str) -> Result<()> {
let node = self
.nodes
.get_mut(node_id)
.ok_or_else(|| Error::Other(format!("session tree has no node `{node_id}`")))?;
node.label = None;
Ok(())
}
pub fn linear_projection(&self) -> Result<Vec<ChatMessage>> {
self.linear_projection_of(&self.active_branch)
}
pub fn linear_projection_of(&self, branch: &str) -> Result<Vec<ChatMessage>> {
let b = self.require_branch(branch)?;
let Some(mut cursor) = b.leaf.clone() else {
return Ok(Vec::new());
};
let mut chain = Vec::new();
let mut visited = BTreeSet::new();
loop {
if !visited.insert(cursor.clone()) {
return Err(Error::Other(format!(
"session tree branch `{branch}` contains a cycle at node `{cursor}`"
)));
}
let node = self.require_node(&cursor)?;
chain.push(node.message.clone());
match &node.parent {
Some(p) => cursor = p.clone(),
None => break,
}
}
chain.reverse();
Ok(chain)
}
pub fn active_path(&self) -> Result<Vec<NodeId>> {
self.path_of(&self.active_branch)
}
pub fn path_of(&self, branch: &str) -> Result<Vec<NodeId>> {
let b = self.require_branch(branch)?;
let Some(mut cursor) = b.leaf.clone() else {
return Ok(Vec::new());
};
let mut chain = Vec::new();
let mut visited = BTreeSet::new();
loop {
if !visited.insert(cursor.clone()) {
return Err(Error::Other(format!(
"session tree branch `{branch}` contains a cycle at node `{cursor}`"
)));
}
let node = self.require_node(&cursor)?;
chain.push(cursor.clone());
match &node.parent {
Some(p) => cursor = p.clone(),
None => break,
}
}
chain.reverse();
Ok(chain)
}
pub fn has_branches(&self) -> bool {
self.branches.len() > 1
}
pub fn summarize_branch(
&mut self,
branch: &str,
summary: impl Into<String>,
model_id: Option<String>,
timestamp_ms: i64,
) -> Result<()> {
let leaf = self.require_branch(branch)?.leaf.clone().ok_or_else(|| {
Error::Other(format!(
"session tree branch `{branch}` has no leaf yet — nothing to summarize"
))
})?;
let b = self
.branches
.get_mut(branch)
.expect("just checked via require_branch");
b.summary = Some(BranchSummary {
summary: summary.into(),
node_id: leaf,
branch: branch.to_string(),
model_id,
created_at_ms: timestamp_ms,
});
Ok(())
}
pub fn render_branch_text(&self, branch: &str) -> Result<String> {
let messages = self.linear_projection_of(branch)?;
let mut out = String::new();
for m in &messages {
let role = match m.role {
crate::message::Role::System => "system",
crate::message::Role::User => "user",
crate::message::Role::Assistant => "assistant",
crate::message::Role::Tool => "tool",
};
out.push_str(role);
out.push_str(": ");
out.push_str(m.content.as_deref().unwrap_or(""));
out.push('\n');
}
Ok(out)
}
pub fn summarize_branch_with(
&mut self,
branch: &str,
summarizer: &dyn BranchSummarizer,
timestamp_ms: i64,
) -> Result<()> {
let text = self.render_branch_text(branch)?;
let turn_count = self.linear_projection_of(branch)?.len();
match summarizer.summarize(&text) {
Ok(summary) => {
self.summarize_branch(
branch,
summary,
Some(summarizer.model_id().to_string()),
timestamp_ms,
)?;
}
Err(_) => {
self.summarize_branch(
branch,
format!("[{turn_count} turn(s), unsummarized]"),
None,
timestamp_ms,
)?;
}
}
Ok(())
}
pub fn splice_for_linear_export(&self) -> Result<(Vec<ChatMessage>, Vec<BranchSummary>)> {
let active = self.linear_projection()?;
let mut summaries = Vec::new();
for (name, b) in &self.branches {
if name == &self.active_branch {
continue;
}
if let Some(s) = &b.summary {
summaries.push(s.clone());
} else {
let (summary_text, node_id) = match &b.leaf {
None => (
format!("[branch `{name}` has no leaf yet — nothing to summarize]"),
String::new(),
),
Some(leaf) => match self.linear_projection_of(name) {
Ok(msgs) => (
format!("[{} turn(s), unsummarized]", msgs.len()),
leaf.clone(),
),
Err(e) => (
format!("[branch `{name}` could not be read, unsummarized: {e}]"),
leaf.clone(),
),
},
};
summaries.push(BranchSummary {
summary: summary_text,
node_id,
branch: name.clone(),
model_id: None,
created_at_ms: b.created_at_ms,
});
}
}
Ok((active, summaries))
}
}
pub trait BranchSummarizer {
fn summarize(&self, branch_text: &str) -> Result<String>;
fn model_id(&self) -> &str;
}