oxi_store/
session_navigation.rs

1//! Session tree navigation for branched sessions
2//!
3//! Handles navigation between different points in the session tree,
4//! including the creation of branch summaries when leaving a path.
5
6use chrono::Utc;
7use serde::{Deserialize, Serialize};
8use std::collections::{HashMap, HashSet};
9use uuid::Uuid;
10
11// ============================================================================
12// Types
13// ============================================================================
14
15/// Extended session entry types for tree navigation
16#[derive(Debug, Clone, Serialize, Deserialize)]
17#[serde(tag = "type")]
18pub enum SessionEntryType {
19    /// A regular message (user/assistant/system/tool)
20    Message(MessageEntry),
21    /// Summary of a branch when navigating away
22    BranchSummary(BranchSummaryEntry),
23    /// Summary after context compaction
24    Compaction(CompactionEntry),
25    /// A label/marker on an entry
26    Label(LabelEntry),
27    /// User-defined metadata for the session
28    SessionInfo(SessionInfoEntry),
29    /// Extension-specific data (not sent to LLM)
30    Custom(CustomEntry),
31    /// Custom message from extension (sent to LLM)
32    CustomMessage(CustomMessageEntry),
33}
34
35/// Message entry content
36#[derive(Debug, Clone, Serialize, Deserialize)]
37pub struct MessageEntry {
38    /// Unique identifier for this message.
39    pub id: Uuid,
40    /// Parent entry ID in the session tree.
41    pub parent_id: Option<Uuid>,
42    /// Unix timestamp in milliseconds.
43    pub timestamp: i64,
44    /// Role of the message sender.
45    pub role: MessageRole,
46    /// Text content of the message.
47    pub content: String,
48}
49
50/// Message role types
51#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
52#[serde(rename_all = "camelCase")]
53pub enum MessageRole {
54    /// User-authored message.
55    User,
56    /// Assistant-authored message.
57    Assistant,
58    /// System prompt message.
59    System,
60    /// Tool call message.
61    Tool,
62    /// Tool result message.
63    ToolResult,
64    /// Custom role message.
65    Custom,
66}
67
68impl MessageRole {
69    /// Check if this is a user message
70    pub fn is_user(&self) -> bool {
71        matches!(self, MessageRole::User)
72    }
73
74    /// Check if this is an assistant message
75    pub fn is_assistant(&self) -> bool {
76        matches!(self, MessageRole::Assistant)
77    }
78}
79
80/// Branch summary entry - captures context when navigating away from a branch
81#[derive(Debug, Clone, Serialize, Deserialize)]
82pub struct BranchSummaryEntry {
83    /// Unique identifier for this summary entry.
84    pub id: Uuid,
85    /// Parent entry ID in the session tree.
86    pub parent_id: Option<Uuid>,
87    /// Unix timestamp in milliseconds.
88    pub timestamp: i64,
89    /// ID of the entry this summary is attached to
90    pub from_id: Uuid,
91    /// The summary text
92    pub summary: String,
93    /// Optional details (e.g., file lists)
94    #[serde(skip_serializing_if = "Option::is_none")]
95    pub details: Option<BranchSummaryDetails>,
96    /// True if generated by an extension
97    #[serde(skip_serializing_if = "Option::is_none")]
98    pub from_hook: Option<bool>,
99}
100
101/// Details stored in BranchSummaryEntry for file tracking
102#[derive(Debug, Clone, Serialize, Deserialize)]
103pub struct BranchSummaryDetails {
104    /// Files that were read during the branch.
105    pub read_files: Vec<String>,
106    /// Files that were modified during the branch.
107    pub modified_files: Vec<String>,
108}
109
110/// Compaction entry - summary after context window compaction
111#[derive(Debug, Clone, Serialize, Deserialize)]
112pub struct CompactionEntry {
113    /// Unique identifier for this compaction entry.
114    pub id: Uuid,
115    /// Parent entry ID in the session tree.
116    pub parent_id: Option<Uuid>,
117    /// Unix timestamp in milliseconds.
118    pub timestamp: i64,
119    /// Compacted summary text.
120    pub summary: String,
121    /// ID of first entry kept after compaction
122    pub first_kept_entry_id: Uuid,
123    /// Estimated tokens before compaction
124    pub tokens_before: usize,
125    /// Extension-specific data
126    #[serde(skip_serializing_if = "Option::is_none")]
127    pub details: Option<serde_json::Value>,
128    /// True if generated by an extension
129    #[serde(skip_serializing_if = "Option::is_none")]
130    pub from_hook: Option<bool>,
131}
132
133/// Label entry - user-defined bookmark/marker on entries
134#[derive(Debug, Clone, Serialize, Deserialize)]
135pub struct LabelEntry {
136    /// Unique identifier for this label entry.
137    pub id: Uuid,
138    /// Parent entry ID in the session tree.
139    pub parent_id: Option<Uuid>,
140    /// Unix timestamp in milliseconds.
141    pub timestamp: i64,
142    /// Entry ID this label is attached to.
143    pub target_id: Uuid,
144    /// Optional label text.
145    pub label: Option<String>,
146}
147
148/// Session info entry - metadata like display name
149#[derive(Debug, Clone, Serialize, Deserialize)]
150pub struct SessionInfoEntry {
151    /// Unique identifier for this session info entry.
152    pub id: Uuid,
153    /// Parent entry ID in the session tree.
154    pub parent_id: Option<Uuid>,
155    /// Unix timestamp in milliseconds.
156    pub timestamp: i64,
157    /// Optional display name for the session.
158    pub name: Option<String>,
159}
160
161/// Custom entry for extensions to store extension-specific data
162#[derive(Debug, Clone, Serialize, Deserialize)]
163pub struct CustomEntry {
164    /// Unique identifier for this custom entry.
165    pub id: Uuid,
166    /// Parent entry ID in the session tree.
167    pub parent_id: Option<Uuid>,
168    /// Unix timestamp in milliseconds.
169    pub timestamp: i64,
170    /// Extension-defined type discriminator.
171    pub custom_type: String,
172    /// Optional extension-specific data payload.
173    #[serde(skip_serializing_if = "Option::is_none")]
174    pub data: Option<serde_json::Value>,
175}
176
177/// Custom message entry for extensions to inject messages into LLM context
178#[derive(Debug, Clone, Serialize, Deserialize)]
179pub struct CustomMessageEntry {
180    /// Unique identifier for this custom message entry.
181    pub id: Uuid,
182    /// Parent entry ID in the session tree.
183    pub parent_id: Option<Uuid>,
184    /// Unix timestamp in milliseconds.
185    pub timestamp: i64,
186    /// Extension-defined type discriminator.
187    pub custom_type: String,
188    /// Message content injected into LLM context.
189    pub content: String,
190    /// Optional extension-specific details.
191    #[serde(skip_serializing_if = "Option::is_none")]
192    pub details: Option<serde_json::Value>,
193    /// Controls TUI rendering
194    pub display: bool,
195}
196
197/// Result of collecting entries for branch summarization
198#[derive(Debug, Clone)]
199pub struct CollectEntriesResult {
200    /// Entries to summarize, in chronological order
201    pub entries: Vec<SessionEntryType>,
202    /// Common ancestor between old and new position
203    pub common_ancestor_id: Option<Uuid>,
204}
205
206/// Navigation target options
207#[derive(Debug, Clone, Default)]
208pub struct NavigationOptions {
209    /// If true, generate a summary of the branch being left
210    pub summarize: bool,
211    /// Custom instructions for summarization
212    pub custom_instructions: Option<String>,
213    /// If true, custom_instructions replaces the default prompt instead of being appended
214    pub replace_instructions: bool,
215    /// Label to attach to the branch summary entry
216    pub label: Option<String>,
217}
218
219/// Result of tree navigation
220#[derive(Debug, Clone)]
221pub struct NavigationResult {
222    /// Text content for the editor (user message text when navigating to user message)
223    pub editor_text: Option<String>,
224    /// True if navigation was cancelled
225    pub cancelled: bool,
226    /// True if summarization was aborted
227    pub aborted: bool,
228    /// ID of the summary entry created, if any
229    pub summary_entry_id: Option<Uuid>,
230}
231
232/// Preparation data for tree navigation (passed to extension hooks)
233#[derive(Debug, Clone)]
234pub struct TreePreparation {
235    /// The navigation target entry ID.
236    pub target_id: Uuid,
237    /// The leaf ID before navigation.
238    pub old_leaf_id: Option<Uuid>,
239    /// Common ancestor between old and new positions.
240    pub common_ancestor_id: Option<Uuid>,
241    /// Entries collected for summarization.
242    pub entries_to_summarize: Vec<SessionEntryType>,
243    /// Whether the user requested a summary.
244    pub user_wants_summary: bool,
245    /// Custom instructions for summarization.
246    pub custom_instructions: Option<String>,
247    /// Whether custom instructions replace the default prompt.
248    pub replace_instructions: bool,
249    /// Label to attach to the summary.
250    pub label: Option<String>,
251}
252
253/// Trait for LLM-based summarization
254/// Implement this to provide actual summarization, or use a mock for testing
255pub trait Summarizer: Send + Sync {
256    /// Generate a summary of the given entries
257    fn summarize(
258        &self,
259        entries: &[SessionEntryType],
260        custom_instructions: Option<&str>,
261        replace_instructions: bool,
262    ) -> impl std::future::Future<Output = Result<BranchSummaryResult, SummarizationError>>
263           + Send
264           + 'static;
265}
266
267/// Result of branch summarization
268#[derive(Debug, Clone)]
269pub struct BranchSummaryResult {
270    /// The generated summary text
271    pub summary: Option<String>,
272    /// Files that were read during this branch
273    pub read_files: Vec<String>,
274    /// Files that were modified during this branch
275    pub modified_files: Vec<String>,
276    /// True if the operation was aborted
277    pub aborted: bool,
278    /// Error message if summarization failed
279    pub error: Option<String>,
280}
281
282/// Error during summarization
283#[derive(Debug, Clone)]
284pub enum SummarizationError {
285    /// No model available for summarization.
286    NoModel,
287    /// Summarization was aborted by the user.
288    Aborted,
289    /// Summarization failed with an error message.
290    Failed(String),
291}
292
293impl std::fmt::Display for SummarizationError {
294    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
295        match self {
296            SummarizationError::NoModel => write!(f, "No model available for summarization"),
297            SummarizationError::Aborted => write!(f, "Summarization was aborted"),
298            SummarizationError::Failed(msg) => write!(f, "Summarization failed: {}", msg),
299        }
300    }
301}
302
303/// Extension hook result for session_before_tree event
304#[derive(Debug, Clone)]
305pub struct BeforeTreeHookResult {
306    /// If true, cancel the navigation
307    pub cancel: bool,
308    /// Summary from extension (if user wants summary)
309    pub summary: Option<ExtensionSummary>,
310    /// Override custom instructions
311    pub custom_instructions: Option<String>,
312    /// Override replace instructions flag
313    pub replace_instructions: Option<bool>,
314    /// Override label
315    pub label: Option<String>,
316}
317
318/// Summary generated by an extension
319#[derive(Debug, Clone)]
320pub struct ExtensionSummary {
321    /// Summary text generated by the extension.
322    pub summary: String,
323    /// Optional extension-specific details.
324    pub details: Option<serde_json::Value>,
325}
326
327// ============================================================================
328// Tree Navigation Implementation
329// ============================================================================
330
331/// Session navigator for tree-based session management
332pub struct SessionNavigator {
333    /// All entries indexed by ID
334    entries_by_id: HashMap<Uuid, SessionEntryType>,
335    /// Labels by entry ID
336    labels_by_id: HashMap<Uuid, String>,
337    /// Label timestamps by entry ID
338    label_timestamps_by_id: HashMap<Uuid, i64>,
339    /// Current leaf ID (where new entries are appended)
340    leaf_id: Option<Uuid>,
341}
342
343impl SessionNavigator {
344    /// Create a new session navigator
345    pub fn new() -> Self {
346        Self {
347            entries_by_id: HashMap::new(),
348            labels_by_id: HashMap::new(),
349            label_timestamps_by_id: HashMap::new(),
350            leaf_id: None,
351        }
352    }
353
354    /// Create from existing entries
355    pub fn from_entries(entries: Vec<SessionEntryType>, leaf_id: Option<Uuid>) -> Self {
356        let mut entries_by_id = HashMap::new();
357        for entry in &entries {
358            let id = Self::entry_id(entry);
359            entries_by_id.insert(id, entry.clone());
360        }
361
362        Self {
363            entries_by_id,
364            labels_by_id: HashMap::new(),
365            label_timestamps_by_id: HashMap::new(),
366            leaf_id,
367        }
368    }
369
370    /// Get the current leaf ID
371    pub fn get_leaf_id(&self) -> Option<Uuid> {
372        self.leaf_id
373    }
374
375    /// Get an entry by ID
376    pub fn get_entry(&self, id: Uuid) -> Option<&SessionEntryType> {
377        self.entries_by_id.get(&id)
378    }
379
380    /// Get the label for an entry, if any
381    pub fn get_label(&self, id: Uuid) -> Option<&str> {
382        self.labels_by_id.get(&id).map(|s| s.as_str())
383    }
384
385    /// Get all entries
386    pub fn get_entries(&self) -> Vec<&SessionEntryType> {
387        self.entries_by_id.values().collect()
388    }
389
390    /// Get the branch path from root to a given entry
391    pub fn get_branch(&self, from_id: Option<Uuid>) -> Vec<&SessionEntryType> {
392        let mut path = Vec::new();
393        let start_id = from_id.or(self.leaf_id);
394
395        let mut current_id = start_id;
396        while let Some(id) = current_id {
397            if let Some(entry) = self.entries_by_id.get(&id) {
398                path.insert(0, entry);
399                current_id = Self::entry_parent_id(entry);
400            } else {
401                break;
402            }
403        }
404
405        path
406    }
407
408    /// Get all direct children of an entry
409    pub fn get_children(&self, parent_id: Uuid) -> Vec<&SessionEntryType> {
410        self.entries_by_id
411            .values()
412            .filter(|entry| Self::entry_parent_id(entry) == Some(parent_id))
413            .collect()
414    }
415
416    /// Collect entries that should be summarized when navigating from one position to another.
417    pub fn collect_entries_for_branch_summary(
418        &self,
419        old_leaf_id: Option<Uuid>,
420        target_id: Uuid,
421    ) -> CollectEntriesResult {
422        let old_leaf_id = match old_leaf_id {
423            Some(id) => id,
424            None => {
425                return CollectEntriesResult {
426                    entries: Vec::new(),
427                    common_ancestor_id: None,
428                };
429            }
430        };
431
432        let old_path_ids: HashSet<Uuid> = self
433            .get_branch(Some(old_leaf_id))
434            .iter()
435            .map(|e| Self::entry_id(e))
436            .collect();
437
438        let target_path = self.get_branch(Some(target_id));
439
440        let mut common_ancestor_id: Option<Uuid> = None;
441        for entry in target_path.iter().rev() {
442            let id = Self::entry_id(entry);
443            if old_path_ids.contains(&id) {
444                common_ancestor_id = Some(id);
445                break;
446            }
447        }
448
449        let mut entries: Vec<&SessionEntryType> = Vec::new();
450        let mut current_id: Option<Uuid> = Some(old_leaf_id);
451
452        while let Some(id) = current_id {
453            if current_id == common_ancestor_id {
454                break;
455            }
456            if let Some(entry) = self.entries_by_id.get(&id) {
457                entries.push(entry);
458                current_id = Self::entry_parent_id(entry);
459            } else {
460                break;
461            }
462        }
463
464        entries.reverse();
465
466        CollectEntriesResult {
467            entries: entries.into_iter().cloned().collect(),
468            common_ancestor_id,
469        }
470    }
471
472    /// Navigate to a target entry in the session tree.
473    pub fn navigate_tree<N: Summarizer + ?Sized>(
474        &mut self,
475        target_id: Uuid,
476        options: NavigationOptions,
477        summarizer: Option<&N>,
478        extension_hook: Option<&dyn Fn(TreePreparation) -> BeforeTreeHookResult>,
479    ) -> NavigationResult {
480        let old_leaf_id = self.leaf_id;
481
482        if Some(target_id) == old_leaf_id {
483            return NavigationResult {
484                editor_text: None,
485                cancelled: false,
486                aborted: false,
487                summary_entry_id: None,
488            };
489        }
490
491        let target_entry = match self.entries_by_id.get(&target_id) {
492            Some(e) => e,
493            None => {
494                return NavigationResult {
495                    editor_text: None,
496                    cancelled: true,
497                    aborted: false,
498                    summary_entry_id: None,
499                };
500            }
501        };
502
503        let collection = self.collect_entries_for_branch_summary(old_leaf_id, target_id);
504
505        let mut custom_instructions = options.custom_instructions.clone();
506        let mut replace_instructions = options.replace_instructions;
507        let mut label = options.label.clone();
508
509        let preparation = TreePreparation {
510            target_id,
511            old_leaf_id,
512            common_ancestor_id: collection.common_ancestor_id,
513            entries_to_summarize: collection.entries.clone(),
514            user_wants_summary: options.summarize,
515            custom_instructions: custom_instructions.clone(),
516            replace_instructions,
517            label: label.clone(),
518        };
519
520        let mut extension_summary: Option<ExtensionSummary> = None;
521        let mut from_extension = false;
522
523        if let Some(hook) = extension_hook {
524            let result = hook(preparation);
525            if result.cancel {
526                return NavigationResult {
527                    editor_text: None,
528                    cancelled: true,
529                    aborted: false,
530                    summary_entry_id: None,
531                };
532            }
533
534            if let Some(ext_sum) = result.summary {
535                extension_summary = Some(ext_sum);
536                from_extension = true;
537            }
538
539            if let Some(ci) = result.custom_instructions {
540                custom_instructions = Some(ci);
541            }
542            if let Some(ri) = result.replace_instructions {
543                replace_instructions = ri;
544            }
545            if let Some(l) = result.label {
546                label = Some(l);
547            }
548        }
549
550        let mut summary_text: Option<String> = None;
551        let mut summary_details: Option<BranchSummaryDetails> = None;
552
553        if options.summarize && !collection.entries.is_empty() && extension_summary.is_none() {
554            if let Some(summarizer) = summarizer {
555                // Use tokio::task::block_in_place to safely call block_on from within
556                // a multi-threaded tokio runtime. Falls back to creating a new runtime
557                // if not in a tokio context.
558                let entries_clone: Vec<SessionEntryType> = collection.entries.clone();
559                let custom_clone = custom_instructions.clone();
560                let result = tokio::task::block_in_place(|| {
561                    tokio::runtime::Handle::current().block_on(summarizer.summarize(
562                        &entries_clone,
563                        custom_clone.as_deref(),
564                        replace_instructions,
565                    ))
566                });
567
568                match result {
569                    Ok(summary_result) => {
570                        if summary_result.aborted {
571                            return NavigationResult {
572                                editor_text: None,
573                                cancelled: true,
574                                aborted: true,
575                                summary_entry_id: None,
576                            };
577                        }
578                        if let Some(err) = summary_result.error {
579                            tracing::warn!("Summarization failed: {}", err);
580                        }
581                        summary_text = summary_result.summary;
582                        if !summary_result.read_files.is_empty()
583                            || !summary_result.modified_files.is_empty()
584                        {
585                            summary_details = Some(BranchSummaryDetails {
586                                read_files: summary_result.read_files,
587                                modified_files: summary_result.modified_files,
588                            });
589                        }
590                    }
591                    Err(e) => {
592                        tracing::warn!("Summarization error: {:?}", e);
593                    }
594                }
595            }
596        } else if let Some(ext_sum) = extension_summary {
597            summary_text = Some(ext_sum.summary);
598            summary_details = ext_sum.details.and_then(|d| serde_json::from_value(d).ok());
599        }
600
601        let (new_leaf_id, editor_text) = Self::determine_leaf_and_editor(target_entry);
602
603        let has_summary = summary_text.is_some();
604        let summary_entry_id = if let Some(text) = summary_text {
605            let summary_id =
606                self.branch_with_summary(new_leaf_id, text, summary_details, from_extension);
607
608            if let Some(l) = &label {
609                self.append_label_change(summary_id, Some(l.clone()));
610            }
611
612            Some(summary_id)
613        } else if new_leaf_id.is_none() {
614            self.reset_leaf();
615            None
616        } else {
617            if let Some(id) = new_leaf_id {
618                self.branch(id);
619            } else {
620                self.reset_leaf();
621            }
622            None
623        };
624
625        let has_label = label.is_some();
626        if has_label && !has_summary {
627            self.append_label_change(target_id, label);
628        }
629
630        NavigationResult {
631            editor_text,
632            cancelled: false,
633            aborted: false,
634            summary_entry_id,
635        }
636    }
637
638    fn determine_leaf_and_editor(entry: &SessionEntryType) -> (Option<Uuid>, Option<String>) {
639        match entry {
640            SessionEntryType::Message(msg) if msg.role.is_user() => {
641                let editor_text = if msg.content.is_empty() {
642                    None
643                } else {
644                    Some(msg.content.clone())
645                };
646                (msg.parent_id, editor_text)
647            }
648            SessionEntryType::CustomMessage(custom) => {
649                let editor_text = if custom.content.is_empty() {
650                    None
651                } else {
652                    Some(custom.content.clone())
653                };
654                (custom.parent_id, editor_text)
655            }
656            _ => {
657                let leaf_id = Self::entry_id(entry);
658                (Some(leaf_id), None)
659            }
660        }
661    }
662
663    /// Switch to a different entry (start a new branch).
664    pub fn branch(&mut self, branch_from_id: Uuid) {
665        if !self.entries_by_id.contains_key(&branch_from_id) {
666            return;
667        }
668        self.leaf_id = Some(branch_from_id);
669    }
670
671    /// Reset the leaf pointer to null (before any entries).
672    pub fn reset_leaf(&mut self) {
673        self.leaf_id = None;
674    }
675
676    /// Start a new branch with a summary of the abandoned path.
677    pub fn branch_with_summary(
678        &mut self,
679        branch_from_id: Option<Uuid>,
680        summary: String,
681        details: Option<BranchSummaryDetails>,
682        from_hook: bool,
683    ) -> Uuid {
684        self.leaf_id = branch_from_id;
685
686        let summary_id = Uuid::new_v4();
687        let entry = SessionEntryType::BranchSummary(BranchSummaryEntry {
688            id: summary_id,
689            parent_id: branch_from_id,
690            timestamp: Utc::now().timestamp_millis(),
691            from_id: branch_from_id.unwrap_or(Uuid::nil()),
692            summary,
693            details,
694            from_hook: Some(from_hook),
695        });
696
697        self.entries_by_id.insert(summary_id, entry.clone());
698        summary_id
699    }
700
701    /// Set or clear a label on an entry.
702    pub fn append_label_change(&mut self, target_id: Uuid, label: Option<String>) -> Uuid {
703        if !self.entries_by_id.contains_key(&target_id) {
704            return Uuid::nil();
705        }
706
707        let label_id = Uuid::new_v4();
708        let entry = SessionEntryType::Label(LabelEntry {
709            id: label_id,
710            parent_id: self.leaf_id,
711            timestamp: Utc::now().timestamp_millis(),
712            target_id,
713            label: label.clone(),
714        });
715
716        self.entries_by_id.insert(label_id, entry);
717
718        if let Some(l) = label {
719            self.labels_by_id.insert(target_id, l);
720            self.label_timestamps_by_id
721                .insert(target_id, Utc::now().timestamp_millis());
722        } else {
723            self.labels_by_id.remove(&target_id);
724            self.label_timestamps_by_id.remove(&target_id);
725        }
726
727        label_id
728    }
729
730    /// Add an entry to the session
731    pub fn add_entry(&mut self, entry: SessionEntryType) {
732        let id = Self::entry_id(&entry);
733        self.entries_by_id.insert(id, entry);
734    }
735
736    /// Get the label timestamp for an entry, if any
737    pub fn get_label_timestamp(&self, id: Uuid) -> Option<i64> {
738        self.label_timestamps_by_id.get(&id).copied()
739    }
740
741    fn entry_id(entry: &SessionEntryType) -> Uuid {
742        match entry {
743            SessionEntryType::Message(e) => e.id,
744            SessionEntryType::BranchSummary(e) => e.id,
745            SessionEntryType::Compaction(e) => e.id,
746            SessionEntryType::Label(e) => e.id,
747            SessionEntryType::SessionInfo(e) => e.id,
748            SessionEntryType::Custom(e) => e.id,
749            SessionEntryType::CustomMessage(e) => e.id,
750        }
751    }
752
753    fn entry_parent_id(entry: &SessionEntryType) -> Option<Uuid> {
754        match entry {
755            SessionEntryType::Message(e) => e.parent_id,
756            SessionEntryType::BranchSummary(e) => e.parent_id,
757            SessionEntryType::Compaction(e) => e.parent_id,
758            SessionEntryType::Label(e) => e.parent_id,
759            SessionEntryType::SessionInfo(e) => e.parent_id,
760            SessionEntryType::Custom(e) => e.parent_id,
761            SessionEntryType::CustomMessage(e) => e.parent_id,
762        }
763    }
764}
765
766impl Default for SessionNavigator {
767    fn default() -> Self {
768        Self::new()
769    }
770}
771
772// ============================================================================
773// Utility Functions
774// ============================================================================
775
776/// Extract text content from user message
777pub fn extract_user_message_text(content: &str) -> String {
778    content.to_string()
779}
780
781/// Extract text content from custom message
782pub fn extract_custom_message_text(content: &str) -> String {
783    content.to_string()
784}
785
786/// Check if an entry is a user message
787pub fn is_user_message(entry: &SessionEntryType) -> bool {
788    matches!(
789        entry,
790        SessionEntryType::Message(msg) if msg.role.is_user()
791    )
792}
793
794/// Check if an entry is a custom message
795pub fn is_custom_message(entry: &SessionEntryType) -> bool {
796    matches!(entry, SessionEntryType::CustomMessage(_))
797}
798
799/// Check if an entry is an assistant message
800pub fn is_assistant_message(entry: &SessionEntryType) -> bool {
801    matches!(
802        entry,
803        SessionEntryType::Message(msg) if msg.role.is_assistant()
804    )
805}
806
807#[cfg(test)]
808mod tests {
809    use super::*;
810
811    /// No-op summarizer for tests that don't need summarization.
812    struct NoOpSummarizer;
813    impl Summarizer for NoOpSummarizer {
814        fn summarize(
815            &self,
816            _entries: &[SessionEntryType],
817            _custom_instructions: Option<&str>,
818            _replace_instructions: bool,
819        ) -> std::pin::Pin<
820            Box<
821                dyn std::future::Future<Output = Result<BranchSummaryResult, SummarizationError>>
822                    + Send
823                    + 'static,
824            >,
825        > {
826            Box::pin(async {
827                Ok(BranchSummaryResult {
828                    summary: None,
829                    read_files: vec![],
830                    modified_files: vec![],
831                    aborted: false,
832                    error: None,
833                })
834            })
835        }
836    }
837
838    fn create_message(
839        id: Uuid,
840        parent_id: Option<Uuid>,
841        role: MessageRole,
842        content: &str,
843    ) -> SessionEntryType {
844        SessionEntryType::Message(MessageEntry {
845            id,
846            parent_id,
847            timestamp: 0,
848            role,
849            content: content.to_string(),
850        })
851    }
852
853    fn entry_id(entry: &SessionEntryType) -> Uuid {
854        match entry {
855            SessionEntryType::Message(e) => e.id,
856            SessionEntryType::BranchSummary(e) => e.id,
857            SessionEntryType::Compaction(e) => e.id,
858            SessionEntryType::Label(e) => e.id,
859            SessionEntryType::SessionInfo(e) => e.id,
860            SessionEntryType::Custom(e) => e.id,
861            SessionEntryType::CustomMessage(e) => e.id,
862        }
863    }
864
865    // ===================================================================
866    // Original tests
867    // ===================================================================
868
869    #[test]
870    fn test_navigate_to_user_message() {
871        let mut nav = SessionNavigator::new();
872
873        let root_id = Uuid::new_v4();
874        let user_id = Uuid::new_v4();
875        let assistant_id = Uuid::new_v4();
876
877        nav.add_entry(create_message(root_id, None, MessageRole::User, "Hello"));
878        nav.add_entry(create_message(
879            user_id,
880            Some(root_id),
881            MessageRole::User,
882            "How are you?",
883        ));
884        nav.add_entry(create_message(
885            assistant_id,
886            Some(user_id),
887            MessageRole::Assistant,
888            "I'm fine",
889        ));
890
891        nav.branch(assistant_id);
892
893        let result = nav.navigate_tree(
894            user_id,
895            NavigationOptions::default(),
896            None as Option<&NoOpSummarizer>,
897            None,
898        );
899
900        assert!(!result.cancelled);
901        assert!(!result.aborted);
902        assert_eq!(result.editor_text, Some("How are you?".to_string()));
903        assert_eq!(nav.get_leaf_id(), Some(root_id));
904    }
905
906    #[test]
907    fn test_navigate_to_assistant_message() {
908        let mut nav = SessionNavigator::new();
909
910        let root_id = Uuid::new_v4();
911        let user_id = Uuid::new_v4();
912        let assistant_id = Uuid::new_v4();
913
914        nav.add_entry(create_message(root_id, None, MessageRole::User, "Hello"));
915        nav.add_entry(create_message(
916            user_id,
917            Some(root_id),
918            MessageRole::User,
919            "How are you?",
920        ));
921        nav.add_entry(create_message(
922            assistant_id,
923            Some(user_id),
924            MessageRole::Assistant,
925            "I'm fine",
926        ));
927
928        nav.branch(assistant_id);
929
930        let result = nav.navigate_tree(
931            assistant_id,
932            NavigationOptions::default(),
933            None as Option<&NoOpSummarizer>,
934            None,
935        );
936
937        assert!(!result.cancelled);
938        assert_eq!(nav.get_leaf_id(), Some(assistant_id));
939    }
940
941    #[test]
942    fn test_noop_navigation() {
943        let mut nav = SessionNavigator::new();
944
945        let entry_id = Uuid::new_v4();
946        nav.add_entry(create_message(entry_id, None, MessageRole::User, "Test"));
947        nav.branch(entry_id);
948
949        let result = nav.navigate_tree(
950            entry_id,
951            NavigationOptions::default(),
952            None as Option<&NoOpSummarizer>,
953            None,
954        );
955
956        assert!(!result.cancelled);
957        assert_eq!(result.editor_text, None);
958    }
959
960    #[test]
961    fn test_collect_entries_for_branch_summary() {
962        let root_id = Uuid::new_v4();
963        let user_id = Uuid::new_v4();
964        let assistant_id = Uuid::new_v4();
965        let branch_user_id = Uuid::new_v4();
966        let branch_assistant_id = Uuid::new_v4();
967
968        let entries = vec![
969            create_message(root_id, None, MessageRole::User, "Root"),
970            create_message(user_id, Some(root_id), MessageRole::User, "User"),
971            create_message(
972                assistant_id,
973                Some(user_id),
974                MessageRole::Assistant,
975                "Assistant",
976            ),
977            create_message(
978                branch_user_id,
979                Some(user_id),
980                MessageRole::User,
981                "Branch User",
982            ),
983            create_message(
984                branch_assistant_id,
985                Some(branch_user_id),
986                MessageRole::Assistant,
987                "Branch Assistant",
988            ),
989        ];
990
991        let nav = SessionNavigator::from_entries(entries, Some(branch_assistant_id));
992
993        let result = nav.collect_entries_for_branch_summary(Some(branch_assistant_id), root_id);
994
995        assert_eq!(result.common_ancestor_id, Some(root_id));
996        assert_eq!(result.entries.len(), 3);
997    }
998
999    #[test]
1000    fn test_label_attachment() {
1001        let mut nav = SessionNavigator::new();
1002
1003        let entry_id = Uuid::new_v4();
1004        nav.add_entry(create_message(entry_id, None, MessageRole::User, "Test"));
1005
1006        let label_id = nav.append_label_change(entry_id, Some("Important".to_string()));
1007
1008        assert_eq!(nav.get_label(entry_id), Some("Important"));
1009        assert!(nav.get_entry(label_id).is_some());
1010
1011        nav.append_label_change(entry_id, None);
1012
1013        assert_eq!(nav.get_label(entry_id), None);
1014    }
1015
1016    #[test]
1017    fn test_branch_with_summary() {
1018        let mut nav = SessionNavigator::new();
1019
1020        let entry_id = Uuid::new_v4();
1021        nav.add_entry(create_message(entry_id, None, MessageRole::User, "Test"));
1022
1023        nav.branch(entry_id);
1024
1025        let summary_id =
1026            nav.branch_with_summary(Some(entry_id), "This is a summary".to_string(), None, false);
1027
1028        assert!(nav.get_entry(summary_id).is_some());
1029        assert_eq!(nav.get_leaf_id(), Some(entry_id));
1030
1031        match nav.get_entry(summary_id) {
1032            Some(SessionEntryType::BranchSummary(e)) => {
1033                assert_eq!(e.summary, "This is a summary");
1034            }
1035            _ => panic!("Expected branch summary entry"),
1036        }
1037    }
1038
1039    // ===================================================================
1040    // Additional comprehensive tests
1041    // ===================================================================
1042
1043    #[test]
1044    fn test_new_navigator_has_no_leaf() {
1045        let nav = SessionNavigator::new();
1046        assert!(nav.get_leaf_id().is_none());
1047    }
1048
1049    #[test]
1050    fn test_default_navigator_has_no_leaf() {
1051        let nav = SessionNavigator::default();
1052        assert!(nav.get_leaf_id().is_none());
1053    }
1054
1055    #[test]
1056    fn test_get_entry_returns_none_for_unknown() {
1057        let nav = SessionNavigator::new();
1058        assert!(nav.get_entry(Uuid::new_v4()).is_none());
1059    }
1060
1061    #[test]
1062    fn test_get_branch_returns_empty_when_no_entries() {
1063        let nav = SessionNavigator::new();
1064        assert!(nav.get_branch(None).is_empty());
1065    }
1066
1067    #[test]
1068    fn test_get_branch_returns_full_path() {
1069        let mut nav = SessionNavigator::new();
1070        let root_id = Uuid::new_v4();
1071        let mid_id = Uuid::new_v4();
1072        let leaf_id = Uuid::new_v4();
1073
1074        nav.add_entry(create_message(root_id, None, MessageRole::User, "Root"));
1075        nav.add_entry(create_message(
1076            mid_id,
1077            Some(root_id),
1078            MessageRole::Assistant,
1079            "Mid",
1080        ));
1081        nav.add_entry(create_message(
1082            leaf_id,
1083            Some(mid_id),
1084            MessageRole::User,
1085            "Leaf",
1086        ));
1087        nav.branch(leaf_id);
1088
1089        let branch = nav.get_branch(None);
1090        assert_eq!(branch.len(), 3);
1091        assert_eq!(entry_id(branch[0]), root_id);
1092        assert_eq!(entry_id(branch[1]), mid_id);
1093        assert_eq!(entry_id(branch[2]), leaf_id);
1094    }
1095
1096    #[test]
1097    fn test_get_children() {
1098        let mut nav = SessionNavigator::new();
1099        let parent_id = Uuid::new_v4();
1100        let child_a = Uuid::new_v4();
1101        let child_b = Uuid::new_v4();
1102
1103        nav.add_entry(create_message(parent_id, None, MessageRole::User, "Parent"));
1104        nav.add_entry(create_message(
1105            child_a,
1106            Some(parent_id),
1107            MessageRole::Assistant,
1108            "A",
1109        ));
1110        nav.add_entry(create_message(
1111            child_b,
1112            Some(parent_id),
1113            MessageRole::Assistant,
1114            "B",
1115        ));
1116
1117        let children = nav.get_children(parent_id);
1118        assert_eq!(children.len(), 2);
1119    }
1120
1121    #[test]
1122    fn test_get_children_of_leaf() {
1123        let mut nav = SessionNavigator::new();
1124        let id = Uuid::new_v4();
1125        nav.add_entry(create_message(id, None, MessageRole::User, "Solo"));
1126
1127        let children = nav.get_children(id);
1128        assert!(children.is_empty());
1129    }
1130
1131    #[test]
1132    fn test_branch_switches_leaf() {
1133        let mut nav = SessionNavigator::new();
1134        let id = Uuid::new_v4();
1135        nav.add_entry(create_message(id, None, MessageRole::User, "Test"));
1136
1137        nav.branch(id);
1138        assert_eq!(nav.get_leaf_id(), Some(id));
1139
1140        nav.reset_leaf();
1141        assert_eq!(nav.get_leaf_id(), None);
1142    }
1143
1144    #[test]
1145    fn test_reset_leaf() {
1146        let mut nav = SessionNavigator::new();
1147        let id = Uuid::new_v4();
1148        nav.add_entry(create_message(id, None, MessageRole::User, "Test"));
1149        nav.branch(id);
1150        assert_eq!(nav.get_leaf_id(), Some(id));
1151
1152        nav.reset_leaf();
1153        assert!(nav.get_leaf_id().is_none());
1154    }
1155
1156    #[test]
1157    fn test_from_entries_preserves_leaf() {
1158        let id1 = Uuid::new_v4();
1159        let id2 = Uuid::new_v4();
1160        let entries = vec![
1161            create_message(id1, None, MessageRole::User, "A"),
1162            create_message(id2, Some(id1), MessageRole::Assistant, "B"),
1163        ];
1164        let nav = SessionNavigator::from_entries(entries, Some(id2));
1165        assert_eq!(nav.get_leaf_id(), Some(id2));
1166        assert!(nav.get_entry(id1).is_some());
1167        assert!(nav.get_entry(id2).is_some());
1168    }
1169
1170    #[test]
1171    fn test_navigate_to_nonexistent_returns_cancelled() {
1172        let mut nav = SessionNavigator::new();
1173        let id = Uuid::new_v4();
1174        nav.add_entry(create_message(id, None, MessageRole::User, "Test"));
1175        nav.branch(id);
1176
1177        let result = nav.navigate_tree(
1178            Uuid::new_v4(),
1179            NavigationOptions::default(),
1180            None as Option<&NoOpSummarizer>,
1181            None,
1182        );
1183        assert!(result.cancelled);
1184    }
1185
1186    #[test]
1187    fn test_navigate_to_root_resets_leaf() {
1188        let mut nav = SessionNavigator::new();
1189        let root_id = Uuid::new_v4();
1190        let child_id = Uuid::new_v4();
1191
1192        nav.add_entry(create_message(root_id, None, MessageRole::User, "Root"));
1193        nav.add_entry(create_message(
1194            child_id,
1195            Some(root_id),
1196            MessageRole::Assistant,
1197            "Child",
1198        ));
1199        nav.branch(child_id);
1200
1201        let result = nav.navigate_tree(
1202            root_id,
1203            NavigationOptions::default(),
1204            None as Option<&NoOpSummarizer>,
1205            None,
1206        );
1207        assert!(!result.cancelled);
1208        assert_eq!(result.editor_text, Some("Root".to_string()));
1209        assert_eq!(nav.get_leaf_id(), None);
1210    }
1211
1212    #[test]
1213    fn test_collect_entries_no_old_leaf() {
1214        let target_id = Uuid::new_v4();
1215        let mut nav = SessionNavigator::new();
1216        nav.add_entry(create_message(target_id, None, MessageRole::User, "T"));
1217
1218        let result = nav.collect_entries_for_branch_summary(None, target_id);
1219        assert!(result.entries.is_empty());
1220        assert_eq!(result.common_ancestor_id, None);
1221    }
1222
1223    #[test]
1224    fn test_collect_entries_same_branch_common_ancestor() {
1225        let root_id = Uuid::new_v4();
1226        let user_id = Uuid::new_v4();
1227        let assistant_id = Uuid::new_v4();
1228
1229        let entries = vec![
1230            create_message(root_id, None, MessageRole::User, "Root"),
1231            create_message(user_id, Some(root_id), MessageRole::User, "User"),
1232            create_message(assistant_id, Some(user_id), MessageRole::Assistant, "Asst"),
1233        ];
1234
1235        let nav = SessionNavigator::from_entries(entries, Some(assistant_id));
1236
1237        let result = nav.collect_entries_for_branch_summary(Some(assistant_id), user_id);
1238        assert_eq!(result.common_ancestor_id, Some(user_id));
1239        assert_eq!(result.entries.len(), 1);
1240    }
1241
1242    #[test]
1243    fn test_label_timestamp() {
1244        let mut nav = SessionNavigator::new();
1245        let id = Uuid::new_v4();
1246        nav.add_entry(create_message(id, None, MessageRole::User, "Test"));
1247
1248        assert!(nav.get_label_timestamp(id).is_none());
1249        nav.append_label_change(id, Some("marker".to_string()));
1250        assert!(nav.get_label_timestamp(id).is_some());
1251    }
1252
1253    #[test]
1254    fn test_label_replace() {
1255        let mut nav = SessionNavigator::new();
1256        let id = Uuid::new_v4();
1257        nav.add_entry(create_message(id, None, MessageRole::User, "Test"));
1258
1259        nav.append_label_change(id, Some("first".to_string()));
1260        assert_eq!(nav.get_label(id), Some("first"));
1261
1262        nav.append_label_change(id, Some("second".to_string()));
1263        assert_eq!(nav.get_label(id), Some("second"));
1264    }
1265
1266    #[test]
1267    fn test_label_nonexistent_entry_returns_nil() {
1268        let mut nav = SessionNavigator::new();
1269        let id = nav.append_label_change(Uuid::new_v4(), Some("ghost".to_string()));
1270        assert_eq!(id, Uuid::nil());
1271    }
1272
1273    #[test]
1274    fn test_branch_with_summary_details() {
1275        let mut nav = SessionNavigator::new();
1276        let id = Uuid::new_v4();
1277        nav.add_entry(create_message(id, None, MessageRole::User, "Test"));
1278        nav.branch(id);
1279
1280        let details = BranchSummaryDetails {
1281            read_files: vec!["a.rs".into()],
1282            modified_files: vec!["b.rs".into()],
1283        };
1284        let summary_id = nav.branch_with_summary(Some(id), "Summary".into(), Some(details), true);
1285
1286        match nav.get_entry(summary_id) {
1287            Some(SessionEntryType::BranchSummary(e)) => {
1288                assert_eq!(e.summary, "Summary");
1289                assert!(e.from_hook.unwrap_or(false));
1290                assert!(e.details.is_some());
1291                let d = e.details.as_ref().unwrap();
1292                assert_eq!(d.read_files, vec!["a.rs"]);
1293                assert_eq!(d.modified_files, vec!["b.rs"]);
1294            }
1295            _ => panic!("Expected branch summary"),
1296        }
1297    }
1298
1299    #[test]
1300    fn test_navigate_with_extension_cancel() {
1301        let mut nav = SessionNavigator::new();
1302        let root_id = Uuid::new_v4();
1303        let child_id = Uuid::new_v4();
1304        nav.add_entry(create_message(root_id, None, MessageRole::User, "R"));
1305        nav.add_entry(create_message(
1306            child_id,
1307            Some(root_id),
1308            MessageRole::Assistant,
1309            "C",
1310        ));
1311        nav.branch(child_id);
1312
1313        let hook = |_: TreePreparation| -> BeforeTreeHookResult {
1314            BeforeTreeHookResult {
1315                cancel: true,
1316                summary: None,
1317                custom_instructions: None,
1318                replace_instructions: None,
1319                label: None,
1320            }
1321        };
1322
1323        let result = nav.navigate_tree(
1324            root_id,
1325            NavigationOptions::default(),
1326            None as Option<&NoOpSummarizer>,
1327            Some(&hook),
1328        );
1329        assert!(result.cancelled);
1330    }
1331
1332    #[test]
1333    fn test_navigate_with_extension_summary() {
1334        let mut nav = SessionNavigator::new();
1335        let root_id = Uuid::new_v4();
1336        let child_id = Uuid::new_v4();
1337        nav.add_entry(create_message(root_id, None, MessageRole::User, "R"));
1338        nav.add_entry(create_message(
1339            child_id,
1340            Some(root_id),
1341            MessageRole::Assistant,
1342            "C",
1343        ));
1344        nav.branch(child_id);
1345
1346        let hook = |_: TreePreparation| -> BeforeTreeHookResult {
1347            BeforeTreeHookResult {
1348                cancel: false,
1349                summary: Some(ExtensionSummary {
1350                    summary: "Ext summary".into(),
1351                    details: None,
1352                }),
1353                custom_instructions: None,
1354                replace_instructions: None,
1355                label: None,
1356            }
1357        };
1358
1359        let result = nav.navigate_tree(
1360            root_id,
1361            NavigationOptions {
1362                summarize: true,
1363                ..Default::default()
1364            },
1365            None as Option<&NoOpSummarizer>,
1366            Some(&hook),
1367        );
1368        assert!(!result.cancelled);
1369        assert!(result.summary_entry_id.is_some());
1370
1371        let sid = result.summary_entry_id.unwrap();
1372        match nav.get_entry(sid) {
1373            Some(SessionEntryType::BranchSummary(e)) => {
1374                assert_eq!(e.summary, "Ext summary");
1375            }
1376            _ => panic!("Expected branch summary from extension"),
1377        }
1378    }
1379
1380    #[test]
1381    fn test_message_role_checks() {
1382        assert!(MessageRole::User.is_user());
1383        assert!(!MessageRole::User.is_assistant());
1384        assert!(MessageRole::Assistant.is_assistant());
1385        assert!(!MessageRole::Assistant.is_user());
1386        assert!(!MessageRole::System.is_user());
1387        assert!(!MessageRole::Tool.is_user());
1388    }
1389
1390    #[test]
1391    fn test_utility_functions() {
1392        let user_entry = create_message(Uuid::new_v4(), None, MessageRole::User, "hi");
1393        let asst_entry = create_message(Uuid::new_v4(), None, MessageRole::Assistant, "yo");
1394        let sys_entry = create_message(Uuid::new_v4(), None, MessageRole::System, "sys");
1395
1396        assert!(is_user_message(&user_entry));
1397        assert!(!is_user_message(&asst_entry));
1398        assert!(is_assistant_message(&asst_entry));
1399        assert!(!is_assistant_message(&user_entry));
1400        assert!(!is_user_message(&sys_entry));
1401        assert!(!is_assistant_message(&sys_entry));
1402    }
1403
1404    #[test]
1405    fn test_session_entry_type_accessors() {
1406        let id = Uuid::new_v4();
1407        let msg = SessionEntryType::Message(MessageEntry {
1408            id,
1409            parent_id: None,
1410            timestamp: 42,
1411            role: MessageRole::User,
1412            content: "test".into(),
1413        });
1414
1415        match &msg {
1416            SessionEntryType::Message(e) => {
1417                assert_eq!(e.id, id);
1418                assert_eq!(e.parent_id, None);
1419                assert_eq!(e.timestamp, 42);
1420                assert_eq!(e.role, MessageRole::User);
1421                assert_eq!(e.content, "test");
1422            }
1423            _ => panic!("Expected Message"),
1424        }
1425    }
1426
1427    #[test]
1428    fn test_get_all_entries() {
1429        let mut nav = SessionNavigator::new();
1430        let id1 = Uuid::new_v4();
1431        let id2 = Uuid::new_v4();
1432        nav.add_entry(create_message(id1, None, MessageRole::User, "A"));
1433        nav.add_entry(create_message(id2, Some(id1), MessageRole::Assistant, "B"));
1434
1435        let all = nav.get_entries();
1436        assert_eq!(all.len(), 2);
1437    }
1438}
oxi_store/session_navigation.rs

oxi_store/
session_navigation.rs
