codewhale_tui/
session_tree.rs

1use crate::models::{ContentBlock, Message};
2use chrono::{DateTime, Utc};
3use serde::{Deserialize, Serialize};
4use std::collections::{HashMap, HashSet};
5pub const CURRENT_JOURNAL_SCHEMA_VERSION: u32 = 1;
6pub type EntryId = String;
7#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Default)]
8pub struct SpawnDepth(pub u32);
9impl SpawnDepth {
10    pub fn next(self) -> Self {
11        Self(self.0.saturating_add(1))
12    }
13}
14#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
15#[serde(tag = "kind", rename_all = "snake_case")]
16pub enum SessionEntryKind {
17    Message {
18        message: Message,
19    },
20    User {
21        text: String,
22    },
23    Assistant {
24        text: String,
25    },
26    Compaction {
27        summary: String,
28        #[serde(default, skip_serializing_if = "Option::is_none")]
29        tokens_before: Option<u64>,
30        #[serde(default, skip_serializing_if = "Option::is_none")]
31        tokens_after: Option<u64>,
32        #[serde(default, skip_serializing_if = "Option::is_none")]
33        model: Option<String>,
34    },
35    BranchSummary {
36        branch_id: String,
37        summary: String,
38        #[serde(default, skip_serializing_if = "Option::is_none")]
39        parent_branch_id: Option<String>,
40    },
41    System {
42        content: String,
43    },
44}
45impl SessionEntryKind {
46    pub fn is_contextual(&self) -> bool {
47        matches!(
48            self,
49            Self::Message { .. } | Self::User { .. } | Self::Assistant { .. }
50        )
51    }
52    pub fn as_message(&self) -> Option<Message> {
53        match self {
54            Self::Message { message } => Some(message.clone()),
55            Self::User { text } => Some(Message {
56                role: "user".into(),
57                content: vec![ContentBlock::Text {
58                    text: text.clone(),
59                    cache_control: None,
60                }],
61            }),
62            Self::Assistant { text } => Some(Message {
63                role: "assistant".into(),
64                content: vec![ContentBlock::Text {
65                    text: text.clone(),
66                    cache_control: None,
67                }],
68            }),
69            Self::Compaction { summary, .. } => Some(Message {
70                role: "system".into(),
71                content: vec![ContentBlock::Text {
72                    text: format!("[compaction summary] {summary}"),
73                    cache_control: None,
74                }],
75            }),
76            Self::BranchSummary { summary, .. } => Some(Message {
77                role: "system".into(),
78                content: vec![ContentBlock::Text {
79                    text: format!("[branch summary] {summary}"),
80                    cache_control: None,
81                }],
82            }),
83            Self::System { content } => Some(Message {
84                role: "system".into(),
85                content: vec![ContentBlock::Text {
86                    text: content.clone(),
87                    cache_control: None,
88                }],
89            }),
90        }
91    }
92}
93#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
94pub struct SessionEntry {
95    pub id: EntryId,
96    #[serde(default, skip_serializing_if = "Option::is_none")]
97    pub parent_id: Option<EntryId>,
98    #[serde(flatten)]
99    pub kind: SessionEntryKind,
100    pub created_at: DateTime<Utc>,
101    #[serde(default)]
102    pub spawn_depth: u32,
103}
104impl SessionEntry {
105    pub fn new(kind: SessionEntryKind, parent_id: Option<EntryId>, spawn_depth: u32) -> Self {
106        Self {
107            id: uuid::Uuid::new_v4().to_string(),
108            parent_id,
109            kind,
110            created_at: Utc::now(),
111            spawn_depth,
112        }
113    }
114    pub fn short_id(&self) -> &str {
115        if self.id.len() >= 8 {
116            &self.id[..8]
117        } else {
118            &self.id
119        }
120    }
121}
122#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Default)]
123pub struct SessionJournal {
124    #[serde(default)]
125    pub entries: Vec<SessionEntry>,
126    #[serde(default, skip_serializing_if = "Option::is_none")]
127    pub leaf_id: Option<EntryId>,
128    #[serde(default = "default_journal_schema_version")]
129    pub schema_version: u32,
130    #[serde(default)]
131    pub spawn_depth: u32,
132}
133fn default_journal_schema_version() -> u32 {
134    CURRENT_JOURNAL_SCHEMA_VERSION
135}
136impl SessionJournal {
137    pub fn new() -> Self {
138        Self {
139            entries: Vec::new(),
140            leaf_id: None,
141            schema_version: CURRENT_JOURNAL_SCHEMA_VERSION,
142            spawn_depth: 0,
143        }
144    }
145    pub fn with_spawn_depth(depth: u32) -> Self {
146        Self {
147            spawn_depth: depth,
148            ..Self::new()
149        }
150    }
151    pub fn append(&mut self, kind: SessionEntryKind) -> EntryId {
152        let entry = SessionEntry::new(kind, self.leaf_id.clone(), self.spawn_depth);
153        let id = entry.id.clone();
154        self.entries.push(entry);
155        self.leaf_id = Some(id.clone());
156        id
157    }
158    pub fn append_message(&mut self, message: Message) -> EntryId {
159        self.append(SessionEntryKind::Message { message })
160    }
161    pub fn append_compaction(
162        &mut self,
163        summary: String,
164        tokens_before: Option<u64>,
165        tokens_after: Option<u64>,
166        model: Option<String>,
167    ) -> EntryId {
168        self.append(SessionEntryKind::Compaction {
169            summary,
170            tokens_before,
171            tokens_after,
172            model,
173        })
174    }
175    pub fn append_branch_summary(
176        &mut self,
177        branch_id: String,
178        summary: String,
179        parent_branch_id: Option<String>,
180    ) -> EntryId {
181        self.append(SessionEntryKind::BranchSummary {
182            branch_id,
183            summary,
184            parent_branch_id,
185        })
186    }
187    pub fn branch_to(&mut self, entry_id: &str) -> Result<(), String> {
188        if self.entries.iter().any(|e| e.id == entry_id) {
189            self.leaf_id = Some(entry_id.to_string());
190            Ok(())
191        } else {
192            Err(format!("entry {entry_id} not found"))
193        }
194    }
195    pub fn fork_from(&self, from_entry_id: Option<&str>) -> Result<Self, String> {
196        let leaf = if let Some(id) = from_entry_id {
197            if !self.entries.iter().any(|e| e.id == id) {
198                return Err(format!("fork source {id} not found"));
199            }
200            Some(id.to_string())
201        } else {
202            self.leaf_id.clone()
203        };
204        Ok(Self {
205            entries: self.entries.clone(),
206            leaf_id: leaf,
207            schema_version: self.schema_version,
208            spawn_depth: self.spawn_depth.saturating_add(1),
209        })
210    }
211    pub fn index(&self) -> HashMap<&str, &SessionEntry> {
212        self.entries.iter().map(|e| (e.id.as_str(), e)).collect()
213    }
214    pub fn children_of(&self, parent_id: Option<&str>) -> Vec<&SessionEntry> {
215        self.entries
216            .iter()
217            .filter(|e| e.parent_id.as_deref() == parent_id)
218            .collect()
219    }
220    pub fn contains(&self, entry_id: &str) -> bool {
221        self.entries.iter().any(|e| e.id == entry_id)
222    }
223    pub fn leaf(&self) -> Option<&SessionEntry> {
224        self.leaf_id
225            .as_deref()
226            .and_then(|id| self.entries.iter().find(|e| e.id == id))
227    }
228    pub fn root_to_leaf(&self) -> Vec<&SessionEntry> {
229        let index: HashMap<&str, &SessionEntry> =
230            self.entries.iter().map(|e| (e.id.as_str(), e)).collect();
231        let mut path = Vec::new();
232        let mut cur = self.leaf_id.as_deref();
233        let mut seen = HashSet::new();
234        while let Some(id) = cur {
235            if !seen.insert(id) {
236                break;
237            }
238            if let Some(entry) = index.get(id) {
239                path.push(*entry);
240                cur = entry.parent_id.as_deref();
241            } else {
242                break;
243            }
244        }
245        path.reverse();
246        path
247    }
248    pub fn active_messages(&self, include_system: bool) -> Vec<Message> {
249        self.root_to_leaf()
250            .into_iter()
251            .filter_map(|e| {
252                if !include_system && !e.kind.is_contextual() {
253                    return None;
254                }
255                e.kind.as_message()
256            })
257            .collect()
258    }
259    pub fn leaves(&self) -> Vec<&SessionEntry> {
260        let parents: HashSet<&str> = self
261            .entries
262            .iter()
263            .filter_map(|e| e.parent_id.as_deref())
264            .collect();
265        self.entries
266            .iter()
267            .filter(|e| !parents.contains(e.id.as_str()))
268            .collect()
269    }
270    pub fn is_empty(&self) -> bool {
271        self.entries.is_empty()
272    }
273    pub fn len(&self) -> usize {
274        self.entries.len()
275    }
276    pub fn validate(&self) -> Result<(), String> {
277        let ids: HashSet<&str> = self.entries.iter().map(|e| e.id.as_str()).collect();
278        for entry in &self.entries {
279            if let Some(parent) = entry.parent_id.as_deref()
280                && !ids.contains(parent)
281            {
282                return Err(format!("entry {} missing parent {}", entry.id, parent));
283            }
284        }
285        if let Some(leaf) = self.leaf_id.as_deref()
286            && !ids.contains(leaf)
287        {
288            return Err(format!("leaf {leaf} not found"));
289        }
290        Ok(())
291    }
292    pub fn from_messages(messages: Vec<Message>, spawn_depth: u32) -> Self {
293        let mut j = Self::with_spawn_depth(spawn_depth);
294        for msg in messages {
295            j.append(SessionEntryKind::Message { message: msg });
296        }
297        j
298    }
299    pub fn to_messages(&self) -> Vec<Message> {
300        self.active_messages(true)
301    }
302
303    /// Make `messages` the active projection without rewriting the journal.
304    ///
305    /// The existing active branch remains as evidence. We reuse its longest
306    /// unchanged prefix, then append the repaired suffix as a sibling branch.
307    pub fn rebranch_active_messages(&mut self, messages: &[Message]) {
308        let active_path = self.root_to_leaf();
309        let shared_prefix = active_path
310            .iter()
311            .zip(messages)
312            .take_while(|(entry, message)| entry.kind.as_message().as_ref() == Some(*message))
313            .count();
314        self.leaf_id = shared_prefix
315            .checked_sub(1)
316            .map(|index| active_path[index].id.clone());
317        for message in &messages[shared_prefix..] {
318            self.append_message(message.clone());
319        }
320    }
321}
322#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
323pub struct SessionImportContainer {
324    pub format_version: u32,
325    pub source: String,
326    pub metadata: Option<serde_json::Value>,
327    pub entries: Vec<SessionEntry>,
328    #[serde(default, skip_serializing_if = "Option::is_none")]
329    pub leaf_id: Option<EntryId>,
330    pub exported_at: DateTime<Utc>,
331    #[serde(default)]
332    pub spawn_depth: u32,
333}
334impl SessionImportContainer {
335    pub fn new(
336        source: String,
337        journal: &SessionJournal,
338        metadata: Option<serde_json::Value>,
339    ) -> Self {
340        Self {
341            format_version: CURRENT_JOURNAL_SCHEMA_VERSION,
342            source,
343            metadata,
344            entries: journal.entries.clone(),
345            leaf_id: journal.leaf_id.clone(),
346            exported_at: Utc::now(),
347            spawn_depth: journal.spawn_depth,
348        }
349    }
350    pub fn into_journal(self) -> Result<SessionJournal, String> {
351        let j = SessionJournal {
352            entries: self.entries,
353            leaf_id: self.leaf_id,
354            schema_version: self.format_version,
355            spawn_depth: self.spawn_depth,
356        };
357        j.validate()?;
358        Ok(j)
359    }
360    pub fn to_json(&self) -> Result<String, serde_json::Error> {
361        serde_json::to_string_pretty(self)
362    }
363    pub fn from_json(json: &str) -> Result<Self, serde_json::Error> {
364        serde_json::from_str(json)
365    }
366}
367pub fn render_tree(journal: &SessionJournal) -> String {
368    if journal.entries.is_empty() {
369        return "(empty session — no entries yet)".to_string();
370    }
371    let index = journal.index();
372    let mut out = String::new();
373    let mut children: HashMap<Option<&str>, Vec<&SessionEntry>> = HashMap::new();
374    for entry in &journal.entries {
375        children
376            .entry(entry.parent_id.as_deref())
377            .or_default()
378            .push(entry);
379    }
380    let active_ids: HashSet<&str> = journal
381        .root_to_leaf()
382        .iter()
383        .map(|e| e.id.as_str())
384        .collect();
385    let leaf = journal.leaf_id.as_deref();
386    fn render_node(
387        out: &mut String,
388        children: &HashMap<Option<&str>, Vec<&SessionEntry>>,
389        active_ids: &HashSet<&str>,
390        leaf: Option<&str>,
391        parent: Option<&str>,
392        depth: usize,
393    ) {
394        let Some(nodes) = children.get(&parent) else {
395            return;
396        };
397        for (idx, entry) in nodes.iter().enumerate() {
398            let is_last = idx + 1 == nodes.len();
399            let prefix = if depth == 0 {
400                "".to_string()
401            } else {
402                let mut p = String::new();
403                for _ in 0..depth - 1 {
404                    p.push_str("│  ");
405                }
406                if is_last {
407                    p.push_str("└─ ");
408                } else {
409                    p.push_str("├─ ");
410                }
411                p
412            };
413            let marker = if Some(entry.id.as_str()) == leaf {
414                "*"
415            } else if active_ids.contains(entry.id.as_str()) {
416                "●"
417            } else {
418                "○"
419            };
420            let kind_label = match &entry.kind {
421                SessionEntryKind::Message { message } => {
422                    let role = &message.role;
423                    let snippet: String = message
424                        .content
425                        .iter()
426                        .filter_map(|b| match b {
427                            ContentBlock::Text { text, .. } => Some(text.as_str()),
428                            _ => None,
429                        })
430                        .collect::<Vec<_>>()
431                        .join(" ");
432                    let short: String = snippet.chars().take(60).collect();
433                    format!("{role}: {short}")
434                }
435                SessionEntryKind::User { text } => {
436                    let short: String = text.chars().take(60).collect();
437                    format!("user: {short}")
438                }
439                SessionEntryKind::Assistant { text } => {
440                    let short: String = text.chars().take(60).collect();
441                    format!("assistant: {short}")
442                }
443                SessionEntryKind::Compaction { summary, .. } => {
444                    let short: String = summary.chars().take(60).collect();
445                    format!("compaction: {short}")
446                }
447                SessionEntryKind::BranchSummary {
448                    branch_id, summary, ..
449                } => {
450                    let short: String = summary.chars().take(60).collect();
451                    format!("branch:{} {short}", &branch_id[..branch_id.len().min(8)])
452                }
453                SessionEntryKind::System { content } => {
454                    let short: String = content.chars().take(60).collect();
455                    format!("system: {short}")
456                }
457            };
458            out.push_str(&format!(
459                "{prefix}{marker} {} [{}] {kind_label}\n",
460                entry.short_id(),
461                entry.id
462            ));
463            render_node(out, children, active_ids, leaf, Some(&entry.id), depth + 1);
464        }
465    }
466    render_node(&mut out, &children, &active_ids, leaf, None, 0);
467    let _ = index;
468    if let Some(leaf_id) = leaf {
469        out.push_str(&format!(
470            "\nleaf: {leaf_id} (active, {} entries)\n",
471            journal.entries.len()
472        ));
473    }
474    out
475}
476#[cfg(test)]
477mod tests {
478    use super::*;
479    use crate::models::{ContentBlock, Message};
480    fn msg(role: &str, text: &str) -> Message {
481        Message {
482            role: role.to_string(),
483            content: vec![ContentBlock::Text {
484                text: text.to_string(),
485                cache_control: None,
486            }],
487        }
488    }
489    #[test]
490    fn append_creates_child_of_leaf() {
491        let mut j = SessionJournal::new();
492        let a = j.append(SessionEntryKind::User {
493            text: "hello".into(),
494        });
495        let b = j.append(SessionEntryKind::Assistant {
496            text: "world".into(),
497        });
498        assert_eq!(j.leaf_id.as_deref(), Some(b.as_str()));
499        let be = j.entries.iter().find(|e| e.id == b).unwrap();
500        assert_eq!(be.parent_id.as_deref(), Some(a.as_str()));
501    }
502    #[test]
503    fn branch_moves_leaf_only() {
504        let mut j = SessionJournal::new();
505        let a = j.append(SessionEntryKind::User { text: "a".into() });
506        let b = j.append(SessionEntryKind::User { text: "b".into() });
507        let _c = j.append(SessionEntryKind::User { text: "c".into() });
508        j.branch_to(&a).unwrap();
509        let d = j.append(SessionEntryKind::User { text: "d".into() });
510        assert_eq!(j.entries.len(), 4);
511        assert_eq!(j.leaf_id.as_deref(), Some(d.as_str()));
512        let path: Vec<String> = j.root_to_leaf().iter().map(|e| e.id.clone()).collect();
513        assert_eq!(path, vec![a.clone(), d.clone()]);
514        assert!(j.entries.iter().any(|e| e.id == b));
515    }
516    #[test]
517    fn from_messages_migrates() {
518        let msgs = vec![msg("user", "hi"), msg("assistant", "hello")];
519        let j = SessionJournal::from_messages(msgs, 0);
520        assert_eq!(j.entries.len(), 2);
521        assert!(j.validate().is_ok());
522        assert_eq!(j.root_to_leaf().len(), 2);
523    }
524    #[test]
525    fn repaired_messages_form_an_append_only_sibling_branch() {
526        let original = vec![
527            msg("user", "shared"),
528            msg("assistant", "broken"),
529            msg("user", "old tail"),
530        ];
531        let mut journal = SessionJournal::from_messages(original, 0);
532        let old_leaf = journal.leaf_id.clone().expect("old leaf");
533        let repaired = vec![
534            msg("user", "shared"),
535            msg("assistant", "repaired"),
536            msg("user", "new tail"),
537        ];
538
539        journal.rebranch_active_messages(&repaired);
540
541        assert_eq!(journal.to_messages(), repaired);
542        assert!(journal.contains(&old_leaf), "old evidence must remain");
543        assert_eq!(
544            journal.entries.len(),
545            5,
546            "one shared entry plus two branches"
547        );
548    }
549    #[test]
550    fn compaction_fits() {
551        let mut j = SessionJournal::new();
552        let id = j.append_compaction("summary".into(), Some(1000), Some(100), None);
553        assert!(j.contains(&id));
554        assert!(matches!(
555            j.leaf().unwrap().kind,
556            SessionEntryKind::Compaction { .. }
557        ));
558    }
559    #[test]
560    fn branch_summary_fits() {
561        let mut j = SessionJournal::new();
562        let a = j.append(SessionEntryKind::User {
563            text: "root".into(),
564        });
565        let b = j.append_branch_summary(a.clone(), "branch summary".into(), None);
566        assert!(j.contains(&b));
567    }
568    #[test]
569    fn spawn_depth_fork() {
570        let mut j = SessionJournal::with_spawn_depth(1);
571        let forked = j.fork_from(None).unwrap();
572        assert_eq!(forked.spawn_depth, 2);
573        let a = j.append(SessionEntryKind::User {
574            text: "root".into(),
575        });
576        let fork2 = j.fork_from(Some(&a)).unwrap();
577        assert_eq!(fork2.spawn_depth, 2);
578    }
579    #[test]
580    fn foreign_roundtrip() {
581        let mut j = SessionJournal::new();
582        j.append(SessionEntryKind::User {
583            text: "hello".into(),
584        });
585        let c = SessionImportContainer::new("codewhale".into(), &j, None);
586        let json = c.to_json().unwrap();
587        let back = SessionImportContainer::from_json(&json).unwrap();
588        let j2 = back.into_journal().unwrap();
589        assert_eq!(j.entries.len(), j2.entries.len());
590    }
591    #[test]
592    fn render_marks_active() {
593        let mut j = SessionJournal::new();
594        j.append(SessionEntryKind::User {
595            text: "root".into(),
596        });
597        let tree = render_tree(&j);
598        assert!(tree.contains('*'));
599    }
600    #[test]
601    fn active_messages_root_to_leaf() {
602        let mut j = SessionJournal::new();
603        j.append(SessionEntryKind::User { text: "a".into() });
604        j.append(SessionEntryKind::User { text: "b".into() });
605        let msgs = j.active_messages(false);
606        assert_eq!(msgs.len(), 2);
607        j.branch_to(&j.entries[0].id.clone()).unwrap();
608        j.append(SessionEntryKind::User { text: "c".into() });
609        let msgs2 = j.active_messages(false);
610        assert_eq!(msgs2.len(), 2);
611    }
612}
codewhale_tui/session_tree.rs

codewhale_tui/
session_tree.rs
