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                // Run summarization outside any existing tokio runtime to
559                // avoid "Cannot start a runtime from within a runtime" panics.
560                // Uses std::thread::scope + new_current_thread which works
561                // regardless of whether we're inside a multi-threaded or
562                // current_thread tokio runtime.
563                let entries_clone: Vec<SessionEntryType> = collection.entries.clone();
564                let custom_clone = custom_instructions.clone();
565                let result = std::thread::scope(|s| {
566                    s.spawn(|| {
567                        let rt = tokio::runtime::Builder::new_current_thread()
568                            .enable_all()
569                            .build()
570                            .expect("failed to build temp runtime");
571                        rt.block_on(summarizer.summarize(
572                            &entries_clone,
573                            custom_clone.as_deref(),
574                            replace_instructions,
575                        ))
576                    })
577                    .join()
578                    .expect("summarization thread panicked")
579                });
580
581                match result {
582                    Ok(summary_result) => {
583                        if summary_result.aborted {
584                            return NavigationResult {
585                                editor_text: None,
586                                cancelled: true,
587                                aborted: true,
588                                summary_entry_id: None,
589                            };
590                        }
591                        if let Some(err) = summary_result.error {
592                            tracing::warn!("Summarization failed: {}", err);
593                        }
594                        summary_text = summary_result.summary;
595                        if !summary_result.read_files.is_empty()
596                            || !summary_result.modified_files.is_empty()
597                        {
598                            summary_details = Some(BranchSummaryDetails {
599                                read_files: summary_result.read_files,
600                                modified_files: summary_result.modified_files,
601                            });
602                        }
603                    }
604                    Err(e) => {
605                        tracing::warn!("Summarization error: {:?}", e);
606                    }
607                }
608            }
609        } else if let Some(ext_sum) = extension_summary {
610            summary_text = Some(ext_sum.summary);
611            summary_details = ext_sum.details.and_then(|d| serde_json::from_value(d).ok());
612        }
613
614        let (new_leaf_id, editor_text) = Self::determine_leaf_and_editor(target_entry);
615
616        let has_summary = summary_text.is_some();
617        let summary_entry_id = if let Some(text) = summary_text {
618            let summary_id =
619                self.branch_with_summary(new_leaf_id, text, summary_details, from_extension);
620
621            if let Some(l) = &label {
622                self.append_label_change(summary_id, Some(l.clone()));
623            }
624
625            Some(summary_id)
626        } else if new_leaf_id.is_none() {
627            self.reset_leaf();
628            None
629        } else {
630            if let Some(id) = new_leaf_id {
631                self.branch(id);
632            } else {
633                self.reset_leaf();
634            }
635            None
636        };
637
638        let has_label = label.is_some();
639        if has_label && !has_summary {
640            self.append_label_change(target_id, label);
641        }
642
643        NavigationResult {
644            editor_text,
645            cancelled: false,
646            aborted: false,
647            summary_entry_id,
648        }
649    }
650
651    fn determine_leaf_and_editor(entry: &SessionEntryType) -> (Option<Uuid>, Option<String>) {
652        match entry {
653            SessionEntryType::Message(msg) if msg.role.is_user() => {
654                let editor_text = if msg.content.is_empty() {
655                    None
656                } else {
657                    Some(msg.content.clone())
658                };
659                (msg.parent_id, editor_text)
660            }
661            SessionEntryType::CustomMessage(custom) => {
662                let editor_text = if custom.content.is_empty() {
663                    None
664                } else {
665                    Some(custom.content.clone())
666                };
667                (custom.parent_id, editor_text)
668            }
669            _ => {
670                let leaf_id = Self::entry_id(entry);
671                (Some(leaf_id), None)
672            }
673        }
674    }
675
676    /// Switch to a different entry (start a new branch).
677    pub fn branch(&mut self, branch_from_id: Uuid) {
678        if !self.entries_by_id.contains_key(&branch_from_id) {
679            return;
680        }
681        self.leaf_id = Some(branch_from_id);
682    }
683
684    /// Reset the leaf pointer to null (before any entries).
685    pub fn reset_leaf(&mut self) {
686        self.leaf_id = None;
687    }
688
689    /// Start a new branch with a summary of the abandoned path.
690    pub fn branch_with_summary(
691        &mut self,
692        branch_from_id: Option<Uuid>,
693        summary: String,
694        details: Option<BranchSummaryDetails>,
695        from_hook: bool,
696    ) -> Uuid {
697        self.leaf_id = branch_from_id;
698
699        let summary_id = Uuid::new_v4();
700        let entry = SessionEntryType::BranchSummary(BranchSummaryEntry {
701            id: summary_id,
702            parent_id: branch_from_id,
703            timestamp: Utc::now().timestamp_millis(),
704            from_id: branch_from_id.unwrap_or(Uuid::nil()),
705            summary,
706            details,
707            from_hook: Some(from_hook),
708        });
709
710        self.entries_by_id.insert(summary_id, entry.clone());
711        summary_id
712    }
713
714    /// Set or clear a label on an entry.
715    pub fn append_label_change(&mut self, target_id: Uuid, label: Option<String>) -> Uuid {
716        if !self.entries_by_id.contains_key(&target_id) {
717            return Uuid::nil();
718        }
719
720        let label_id = Uuid::new_v4();
721        let entry = SessionEntryType::Label(LabelEntry {
722            id: label_id,
723            parent_id: self.leaf_id,
724            timestamp: Utc::now().timestamp_millis(),
725            target_id,
726            label: label.clone(),
727        });
728
729        self.entries_by_id.insert(label_id, entry);
730
731        if let Some(l) = label {
732            self.labels_by_id.insert(target_id, l);
733            self.label_timestamps_by_id
734                .insert(target_id, Utc::now().timestamp_millis());
735        } else {
736            self.labels_by_id.remove(&target_id);
737            self.label_timestamps_by_id.remove(&target_id);
738        }
739
740        label_id
741    }
742
743    /// Add an entry to the session
744    pub fn add_entry(&mut self, entry: SessionEntryType) {
745        let id = Self::entry_id(&entry);
746        self.entries_by_id.insert(id, entry);
747    }
748
749    /// Get the label timestamp for an entry, if any
750    pub fn get_label_timestamp(&self, id: Uuid) -> Option<i64> {
751        self.label_timestamps_by_id.get(&id).copied()
752    }
753
754    fn entry_id(entry: &SessionEntryType) -> Uuid {
755        match entry {
756            SessionEntryType::Message(e) => e.id,
757            SessionEntryType::BranchSummary(e) => e.id,
758            SessionEntryType::Compaction(e) => e.id,
759            SessionEntryType::Label(e) => e.id,
760            SessionEntryType::SessionInfo(e) => e.id,
761            SessionEntryType::Custom(e) => e.id,
762            SessionEntryType::CustomMessage(e) => e.id,
763        }
764    }
765
766    fn entry_parent_id(entry: &SessionEntryType) -> Option<Uuid> {
767        match entry {
768            SessionEntryType::Message(e) => e.parent_id,
769            SessionEntryType::BranchSummary(e) => e.parent_id,
770            SessionEntryType::Compaction(e) => e.parent_id,
771            SessionEntryType::Label(e) => e.parent_id,
772            SessionEntryType::SessionInfo(e) => e.parent_id,
773            SessionEntryType::Custom(e) => e.parent_id,
774            SessionEntryType::CustomMessage(e) => e.parent_id,
775        }
776    }
777}
778
779impl Default for SessionNavigator {
780    fn default() -> Self {
781        Self::new()
782    }
783}
784
785// ============================================================================
786// Utility Functions
787// ============================================================================
788
789/// Extract text content from user message
790pub fn extract_user_message_text(content: &str) -> String {
791    content.to_string()
792}
793
794/// Extract text content from custom message
795pub fn extract_custom_message_text(content: &str) -> String {
796    content.to_string()
797}
798
799/// Check if an entry is a user message
800pub fn is_user_message(entry: &SessionEntryType) -> bool {
801    matches!(
802        entry,
803        SessionEntryType::Message(msg) if msg.role.is_user()
804    )
805}
806
807/// Check if an entry is a custom message
808pub fn is_custom_message(entry: &SessionEntryType) -> bool {
809    matches!(entry, SessionEntryType::CustomMessage(_))
810}
811
812/// Check if an entry is an assistant message
813pub fn is_assistant_message(entry: &SessionEntryType) -> bool {
814    matches!(
815        entry,
816        SessionEntryType::Message(msg) if msg.role.is_assistant()
817    )
818}
819
820#[cfg(test)]
821mod tests {
822    use super::*;
823
824    /// No-op summarizer for tests that don't need summarization.
825    struct NoOpSummarizer;
826    impl Summarizer for NoOpSummarizer {
827        fn summarize(
828            &self,
829            _entries: &[SessionEntryType],
830            _custom_instructions: Option<&str>,
831            _replace_instructions: bool,
832        ) -> std::pin::Pin<
833            Box<
834                dyn std::future::Future<Output = Result<BranchSummaryResult, SummarizationError>>
835                    + Send
836                    + 'static,
837            >,
838        > {
839            Box::pin(async {
840                Ok(BranchSummaryResult {
841                    summary: None,
842                    read_files: vec![],
843                    modified_files: vec![],
844                    aborted: false,
845                    error: None,
846                })
847            })
848        }
849    }
850
851    fn create_message(
852        id: Uuid,
853        parent_id: Option<Uuid>,
854        role: MessageRole,
855        content: &str,
856    ) -> SessionEntryType {
857        SessionEntryType::Message(MessageEntry {
858            id,
859            parent_id,
860            timestamp: 0,
861            role,
862            content: content.to_string(),
863        })
864    }
865
866    fn entry_id(entry: &SessionEntryType) -> Uuid {
867        match entry {
868            SessionEntryType::Message(e) => e.id,
869            SessionEntryType::BranchSummary(e) => e.id,
870            SessionEntryType::Compaction(e) => e.id,
871            SessionEntryType::Label(e) => e.id,
872            SessionEntryType::SessionInfo(e) => e.id,
873            SessionEntryType::Custom(e) => e.id,
874            SessionEntryType::CustomMessage(e) => e.id,
875        }
876    }
877
878    // ===================================================================
879    // Original tests
880    // ===================================================================
881
882    #[test]
883    fn test_navigate_to_user_message() {
884        let mut nav = SessionNavigator::new();
885
886        let root_id = Uuid::new_v4();
887        let user_id = Uuid::new_v4();
888        let assistant_id = Uuid::new_v4();
889
890        nav.add_entry(create_message(root_id, None, MessageRole::User, "Hello"));
891        nav.add_entry(create_message(
892            user_id,
893            Some(root_id),
894            MessageRole::User,
895            "How are you?",
896        ));
897        nav.add_entry(create_message(
898            assistant_id,
899            Some(user_id),
900            MessageRole::Assistant,
901            "I'm fine",
902        ));
903
904        nav.branch(assistant_id);
905
906        let result = nav.navigate_tree(
907            user_id,
908            NavigationOptions::default(),
909            None as Option<&NoOpSummarizer>,
910            None,
911        );
912
913        assert!(!result.cancelled);
914        assert!(!result.aborted);
915        assert_eq!(result.editor_text, Some("How are you?".to_string()));
916        assert_eq!(nav.get_leaf_id(), Some(root_id));
917    }
918
919    #[test]
920    fn test_navigate_to_assistant_message() {
921        let mut nav = SessionNavigator::new();
922
923        let root_id = Uuid::new_v4();
924        let user_id = Uuid::new_v4();
925        let assistant_id = Uuid::new_v4();
926
927        nav.add_entry(create_message(root_id, None, MessageRole::User, "Hello"));
928        nav.add_entry(create_message(
929            user_id,
930            Some(root_id),
931            MessageRole::User,
932            "How are you?",
933        ));
934        nav.add_entry(create_message(
935            assistant_id,
936            Some(user_id),
937            MessageRole::Assistant,
938            "I'm fine",
939        ));
940
941        nav.branch(assistant_id);
942
943        let result = nav.navigate_tree(
944            assistant_id,
945            NavigationOptions::default(),
946            None as Option<&NoOpSummarizer>,
947            None,
948        );
949
950        assert!(!result.cancelled);
951        assert_eq!(nav.get_leaf_id(), Some(assistant_id));
952    }
953
954    #[test]
955    fn test_noop_navigation() {
956        let mut nav = SessionNavigator::new();
957
958        let entry_id = Uuid::new_v4();
959        nav.add_entry(create_message(entry_id, None, MessageRole::User, "Test"));
960        nav.branch(entry_id);
961
962        let result = nav.navigate_tree(
963            entry_id,
964            NavigationOptions::default(),
965            None as Option<&NoOpSummarizer>,
966            None,
967        );
968
969        assert!(!result.cancelled);
970        assert_eq!(result.editor_text, None);
971    }
972
973    #[test]
974    fn test_collect_entries_for_branch_summary() {
975        let root_id = Uuid::new_v4();
976        let user_id = Uuid::new_v4();
977        let assistant_id = Uuid::new_v4();
978        let branch_user_id = Uuid::new_v4();
979        let branch_assistant_id = Uuid::new_v4();
980
981        let entries = vec![
982            create_message(root_id, None, MessageRole::User, "Root"),
983            create_message(user_id, Some(root_id), MessageRole::User, "User"),
984            create_message(
985                assistant_id,
986                Some(user_id),
987                MessageRole::Assistant,
988                "Assistant",
989            ),
990            create_message(
991                branch_user_id,
992                Some(user_id),
993                MessageRole::User,
994                "Branch User",
995            ),
996            create_message(
997                branch_assistant_id,
998                Some(branch_user_id),
999                MessageRole::Assistant,
1000                "Branch Assistant",
1001            ),
1002        ];
1003
1004        let nav = SessionNavigator::from_entries(entries, Some(branch_assistant_id));
1005
1006        let result = nav.collect_entries_for_branch_summary(Some(branch_assistant_id), root_id);
1007
1008        assert_eq!(result.common_ancestor_id, Some(root_id));
1009        assert_eq!(result.entries.len(), 3);
1010    }
1011
1012    #[test]
1013    fn test_label_attachment() {
1014        let mut nav = SessionNavigator::new();
1015
1016        let entry_id = Uuid::new_v4();
1017        nav.add_entry(create_message(entry_id, None, MessageRole::User, "Test"));
1018
1019        let label_id = nav.append_label_change(entry_id, Some("Important".to_string()));
1020
1021        assert_eq!(nav.get_label(entry_id), Some("Important"));
1022        assert!(nav.get_entry(label_id).is_some());
1023
1024        nav.append_label_change(entry_id, None);
1025
1026        assert_eq!(nav.get_label(entry_id), None);
1027    }
1028
1029    #[test]
1030    fn test_branch_with_summary() {
1031        let mut nav = SessionNavigator::new();
1032
1033        let entry_id = Uuid::new_v4();
1034        nav.add_entry(create_message(entry_id, None, MessageRole::User, "Test"));
1035
1036        nav.branch(entry_id);
1037
1038        let summary_id =
1039            nav.branch_with_summary(Some(entry_id), "This is a summary".to_string(), None, false);
1040
1041        assert!(nav.get_entry(summary_id).is_some());
1042        assert_eq!(nav.get_leaf_id(), Some(entry_id));
1043
1044        match nav.get_entry(summary_id) {
1045            Some(SessionEntryType::BranchSummary(e)) => {
1046                assert_eq!(e.summary, "This is a summary");
1047            }
1048            _ => panic!("Expected branch summary entry"),
1049        }
1050    }
1051
1052    // ===================================================================
1053    // Additional comprehensive tests
1054    // ===================================================================
1055
1056    #[test]
1057    fn test_new_navigator_has_no_leaf() {
1058        let nav = SessionNavigator::new();
1059        assert!(nav.get_leaf_id().is_none());
1060    }
1061
1062    #[test]
1063    fn test_default_navigator_has_no_leaf() {
1064        let nav = SessionNavigator::default();
1065        assert!(nav.get_leaf_id().is_none());
1066    }
1067
1068    #[test]
1069    fn test_get_entry_returns_none_for_unknown() {
1070        let nav = SessionNavigator::new();
1071        assert!(nav.get_entry(Uuid::new_v4()).is_none());
1072    }
1073
1074    #[test]
1075    fn test_get_branch_returns_empty_when_no_entries() {
1076        let nav = SessionNavigator::new();
1077        assert!(nav.get_branch(None).is_empty());
1078    }
1079
1080    #[test]
1081    fn test_get_branch_returns_full_path() {
1082        let mut nav = SessionNavigator::new();
1083        let root_id = Uuid::new_v4();
1084        let mid_id = Uuid::new_v4();
1085        let leaf_id = Uuid::new_v4();
1086
1087        nav.add_entry(create_message(root_id, None, MessageRole::User, "Root"));
1088        nav.add_entry(create_message(
1089            mid_id,
1090            Some(root_id),
1091            MessageRole::Assistant,
1092            "Mid",
1093        ));
1094        nav.add_entry(create_message(
1095            leaf_id,
1096            Some(mid_id),
1097            MessageRole::User,
1098            "Leaf",
1099        ));
1100        nav.branch(leaf_id);
1101
1102        let branch = nav.get_branch(None);
1103        assert_eq!(branch.len(), 3);
1104        assert_eq!(entry_id(branch[0]), root_id);
1105        assert_eq!(entry_id(branch[1]), mid_id);
1106        assert_eq!(entry_id(branch[2]), leaf_id);
1107    }
1108
1109    #[test]
1110    fn test_get_children() {
1111        let mut nav = SessionNavigator::new();
1112        let parent_id = Uuid::new_v4();
1113        let child_a = Uuid::new_v4();
1114        let child_b = Uuid::new_v4();
1115
1116        nav.add_entry(create_message(parent_id, None, MessageRole::User, "Parent"));
1117        nav.add_entry(create_message(
1118            child_a,
1119            Some(parent_id),
1120            MessageRole::Assistant,
1121            "A",
1122        ));
1123        nav.add_entry(create_message(
1124            child_b,
1125            Some(parent_id),
1126            MessageRole::Assistant,
1127            "B",
1128        ));
1129
1130        let children = nav.get_children(parent_id);
1131        assert_eq!(children.len(), 2);
1132    }
1133
1134    #[test]
1135    fn test_get_children_of_leaf() {
1136        let mut nav = SessionNavigator::new();
1137        let id = Uuid::new_v4();
1138        nav.add_entry(create_message(id, None, MessageRole::User, "Solo"));
1139
1140        let children = nav.get_children(id);
1141        assert!(children.is_empty());
1142    }
1143
1144    #[test]
1145    fn test_branch_switches_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
1150        nav.branch(id);
1151        assert_eq!(nav.get_leaf_id(), Some(id));
1152
1153        nav.reset_leaf();
1154        assert_eq!(nav.get_leaf_id(), None);
1155    }
1156
1157    #[test]
1158    fn test_reset_leaf() {
1159        let mut nav = SessionNavigator::new();
1160        let id = Uuid::new_v4();
1161        nav.add_entry(create_message(id, None, MessageRole::User, "Test"));
1162        nav.branch(id);
1163        assert_eq!(nav.get_leaf_id(), Some(id));
1164
1165        nav.reset_leaf();
1166        assert!(nav.get_leaf_id().is_none());
1167    }
1168
1169    #[test]
1170    fn test_from_entries_preserves_leaf() {
1171        let id1 = Uuid::new_v4();
1172        let id2 = Uuid::new_v4();
1173        let entries = vec![
1174            create_message(id1, None, MessageRole::User, "A"),
1175            create_message(id2, Some(id1), MessageRole::Assistant, "B"),
1176        ];
1177        let nav = SessionNavigator::from_entries(entries, Some(id2));
1178        assert_eq!(nav.get_leaf_id(), Some(id2));
1179        assert!(nav.get_entry(id1).is_some());
1180        assert!(nav.get_entry(id2).is_some());
1181    }
1182
1183    #[test]
1184    fn test_navigate_to_nonexistent_returns_cancelled() {
1185        let mut nav = SessionNavigator::new();
1186        let id = Uuid::new_v4();
1187        nav.add_entry(create_message(id, None, MessageRole::User, "Test"));
1188        nav.branch(id);
1189
1190        let result = nav.navigate_tree(
1191            Uuid::new_v4(),
1192            NavigationOptions::default(),
1193            None as Option<&NoOpSummarizer>,
1194            None,
1195        );
1196        assert!(result.cancelled);
1197    }
1198
1199    #[test]
1200    fn test_navigate_to_root_resets_leaf() {
1201        let mut nav = SessionNavigator::new();
1202        let root_id = Uuid::new_v4();
1203        let child_id = Uuid::new_v4();
1204
1205        nav.add_entry(create_message(root_id, None, MessageRole::User, "Root"));
1206        nav.add_entry(create_message(
1207            child_id,
1208            Some(root_id),
1209            MessageRole::Assistant,
1210            "Child",
1211        ));
1212        nav.branch(child_id);
1213
1214        let result = nav.navigate_tree(
1215            root_id,
1216            NavigationOptions::default(),
1217            None as Option<&NoOpSummarizer>,
1218            None,
1219        );
1220        assert!(!result.cancelled);
1221        assert_eq!(result.editor_text, Some("Root".to_string()));
1222        assert_eq!(nav.get_leaf_id(), None);
1223    }
1224
1225    #[test]
1226    fn test_collect_entries_no_old_leaf() {
1227        let target_id = Uuid::new_v4();
1228        let mut nav = SessionNavigator::new();
1229        nav.add_entry(create_message(target_id, None, MessageRole::User, "T"));
1230
1231        let result = nav.collect_entries_for_branch_summary(None, target_id);
1232        assert!(result.entries.is_empty());
1233        assert_eq!(result.common_ancestor_id, None);
1234    }
1235
1236    #[test]
1237    fn test_collect_entries_same_branch_common_ancestor() {
1238        let root_id = Uuid::new_v4();
1239        let user_id = Uuid::new_v4();
1240        let assistant_id = Uuid::new_v4();
1241
1242        let entries = vec![
1243            create_message(root_id, None, MessageRole::User, "Root"),
1244            create_message(user_id, Some(root_id), MessageRole::User, "User"),
1245            create_message(assistant_id, Some(user_id), MessageRole::Assistant, "Asst"),
1246        ];
1247
1248        let nav = SessionNavigator::from_entries(entries, Some(assistant_id));
1249
1250        let result = nav.collect_entries_for_branch_summary(Some(assistant_id), user_id);
1251        assert_eq!(result.common_ancestor_id, Some(user_id));
1252        assert_eq!(result.entries.len(), 1);
1253    }
1254
1255    #[test]
1256    fn test_label_timestamp() {
1257        let mut nav = SessionNavigator::new();
1258        let id = Uuid::new_v4();
1259        nav.add_entry(create_message(id, None, MessageRole::User, "Test"));
1260
1261        assert!(nav.get_label_timestamp(id).is_none());
1262        nav.append_label_change(id, Some("marker".to_string()));
1263        assert!(nav.get_label_timestamp(id).is_some());
1264    }
1265
1266    #[test]
1267    fn test_label_replace() {
1268        let mut nav = SessionNavigator::new();
1269        let id = Uuid::new_v4();
1270        nav.add_entry(create_message(id, None, MessageRole::User, "Test"));
1271
1272        nav.append_label_change(id, Some("first".to_string()));
1273        assert_eq!(nav.get_label(id), Some("first"));
1274
1275        nav.append_label_change(id, Some("second".to_string()));
1276        assert_eq!(nav.get_label(id), Some("second"));
1277    }
1278
1279    #[test]
1280    fn test_label_nonexistent_entry_returns_nil() {
1281        let mut nav = SessionNavigator::new();
1282        let id = nav.append_label_change(Uuid::new_v4(), Some("ghost".to_string()));
1283        assert_eq!(id, Uuid::nil());
1284    }
1285
1286    #[test]
1287    fn test_branch_with_summary_details() {
1288        let mut nav = SessionNavigator::new();
1289        let id = Uuid::new_v4();
1290        nav.add_entry(create_message(id, None, MessageRole::User, "Test"));
1291        nav.branch(id);
1292
1293        let details = BranchSummaryDetails {
1294            read_files: vec!["a.rs".into()],
1295            modified_files: vec!["b.rs".into()],
1296        };
1297        let summary_id = nav.branch_with_summary(Some(id), "Summary".into(), Some(details), true);
1298
1299        match nav.get_entry(summary_id) {
1300            Some(SessionEntryType::BranchSummary(e)) => {
1301                assert_eq!(e.summary, "Summary");
1302                assert!(e.from_hook.unwrap_or(false));
1303                assert!(e.details.is_some());
1304                let d = e.details.as_ref().unwrap();
1305                assert_eq!(d.read_files, vec!["a.rs"]);
1306                assert_eq!(d.modified_files, vec!["b.rs"]);
1307            }
1308            _ => panic!("Expected branch summary"),
1309        }
1310    }
1311
1312    #[test]
1313    fn test_navigate_with_extension_cancel() {
1314        let mut nav = SessionNavigator::new();
1315        let root_id = Uuid::new_v4();
1316        let child_id = Uuid::new_v4();
1317        nav.add_entry(create_message(root_id, None, MessageRole::User, "R"));
1318        nav.add_entry(create_message(
1319            child_id,
1320            Some(root_id),
1321            MessageRole::Assistant,
1322            "C",
1323        ));
1324        nav.branch(child_id);
1325
1326        let hook = |_: TreePreparation| -> BeforeTreeHookResult {
1327            BeforeTreeHookResult {
1328                cancel: true,
1329                summary: None,
1330                custom_instructions: None,
1331                replace_instructions: None,
1332                label: None,
1333            }
1334        };
1335
1336        let result = nav.navigate_tree(
1337            root_id,
1338            NavigationOptions::default(),
1339            None as Option<&NoOpSummarizer>,
1340            Some(&hook),
1341        );
1342        assert!(result.cancelled);
1343    }
1344
1345    #[test]
1346    fn test_navigate_with_extension_summary() {
1347        let mut nav = SessionNavigator::new();
1348        let root_id = Uuid::new_v4();
1349        let child_id = Uuid::new_v4();
1350        nav.add_entry(create_message(root_id, None, MessageRole::User, "R"));
1351        nav.add_entry(create_message(
1352            child_id,
1353            Some(root_id),
1354            MessageRole::Assistant,
1355            "C",
1356        ));
1357        nav.branch(child_id);
1358
1359        let hook = |_: TreePreparation| -> BeforeTreeHookResult {
1360            BeforeTreeHookResult {
1361                cancel: false,
1362                summary: Some(ExtensionSummary {
1363                    summary: "Ext summary".into(),
1364                    details: None,
1365                }),
1366                custom_instructions: None,
1367                replace_instructions: None,
1368                label: None,
1369            }
1370        };
1371
1372        let result = nav.navigate_tree(
1373            root_id,
1374            NavigationOptions {
1375                summarize: true,
1376                ..Default::default()
1377            },
1378            None as Option<&NoOpSummarizer>,
1379            Some(&hook),
1380        );
1381        assert!(!result.cancelled);
1382        assert!(result.summary_entry_id.is_some());
1383
1384        let sid = result.summary_entry_id.unwrap();
1385        match nav.get_entry(sid) {
1386            Some(SessionEntryType::BranchSummary(e)) => {
1387                assert_eq!(e.summary, "Ext summary");
1388            }
1389            _ => panic!("Expected branch summary from extension"),
1390        }
1391    }
1392
1393    #[test]
1394    fn test_message_role_checks() {
1395        assert!(MessageRole::User.is_user());
1396        assert!(!MessageRole::User.is_assistant());
1397        assert!(MessageRole::Assistant.is_assistant());
1398        assert!(!MessageRole::Assistant.is_user());
1399        assert!(!MessageRole::System.is_user());
1400        assert!(!MessageRole::Tool.is_user());
1401    }
1402
1403    #[test]
1404    fn test_utility_functions() {
1405        let user_entry = create_message(Uuid::new_v4(), None, MessageRole::User, "hi");
1406        let asst_entry = create_message(Uuid::new_v4(), None, MessageRole::Assistant, "yo");
1407        let sys_entry = create_message(Uuid::new_v4(), None, MessageRole::System, "sys");
1408
1409        assert!(is_user_message(&user_entry));
1410        assert!(!is_user_message(&asst_entry));
1411        assert!(is_assistant_message(&asst_entry));
1412        assert!(!is_assistant_message(&user_entry));
1413        assert!(!is_user_message(&sys_entry));
1414        assert!(!is_assistant_message(&sys_entry));
1415    }
1416
1417    #[test]
1418    fn test_session_entry_type_accessors() {
1419        let id = Uuid::new_v4();
1420        let msg = SessionEntryType::Message(MessageEntry {
1421            id,
1422            parent_id: None,
1423            timestamp: 42,
1424            role: MessageRole::User,
1425            content: "test".into(),
1426        });
1427
1428        match &msg {
1429            SessionEntryType::Message(e) => {
1430                assert_eq!(e.id, id);
1431                assert_eq!(e.parent_id, None);
1432                assert_eq!(e.timestamp, 42);
1433                assert_eq!(e.role, MessageRole::User);
1434                assert_eq!(e.content, "test");
1435            }
1436            _ => panic!("Expected Message"),
1437        }
1438    }
1439
1440    #[test]
1441    fn test_get_all_entries() {
1442        let mut nav = SessionNavigator::new();
1443        let id1 = Uuid::new_v4();
1444        let id2 = Uuid::new_v4();
1445        nav.add_entry(create_message(id1, None, MessageRole::User, "A"));
1446        nav.add_entry(create_message(id2, Some(id1), MessageRole::Assistant, "B"));
1447
1448        let all = nav.get_entries();
1449        assert_eq!(all.len(), 2);
1450    }
1451}
oxi_store/session_navigation.rs

oxi_store/
session_navigation.rs
