rab/agent/session/
storage.rs

1//! Low-level session persistence abstraction — Pi-compatible `SessionStorage`.
2//!
3//! Pi architecture:
4//!   SessionStorage (trait) ← InMemorySessionStorage / JsonlSessionStorage
5//!   Session (struct)  ← wraps SessionStorage, provides high-level API
6//!   AgentHarness     ← owns Session, drives agent loop
7//!
8//! This module provides the trait and both implementations.
9//! The `Session` struct lives in `session.rs`.
10
11use super::model::{
12    SessionEntry, SessionHeader, append_entry_to_file, generate_entry_id, load_session_from_file,
13};
14use std::path::{Path, PathBuf};
15
16// ── SessionMetadata ────────────────────────────────────────────────
17
18/// Metadata about a session, derived from the session header.
19/// Pi-compatible: wraps header info into a metadata object.
20#[derive(Debug, Clone)]
21pub struct SessionMetadata {
22    pub id: String,
23    pub created_at: String,
24    pub cwd: String,
25    /// File path on disk, if this is a persisted session.
26    pub path: Option<PathBuf>,
27    /// Path to the parent session if this was forked.
28    pub parent_session_path: Option<String>,
29}
30
31// ── SessionStorage trait ───────────────────────────────────────────
32
33/// Low-level CRUD abstraction for session persistence.
34///
35/// Pi-compatible: provides leaf management, label tracking, path queries,
36/// and entry CRUD. `Session` builds on this for the high-level API.
37pub trait SessionStorage: Send {
38    /// Return header-derived metadata.
39    fn metadata(&self) -> SessionMetadata;
40
41    /// Get the current leaf entry ID (the last non-leaf entry, resolved through leaf entries).
42    /// Returns `None` if no entries exist.
43    fn get_leaf_id(&self) -> Option<String>;
44
45    /// Persist a leaf entry that records the active session-tree leaf.
46    /// `None` means reset to no leaf.
47    fn set_leaf_id(&mut self, leaf_id: Option<&str>) -> Result<(), String>;
48
49    /// Generate a unique 8-character hex entry ID, collision-checked.
50    fn create_entry_id(&self) -> String;
51
52    /// Append a fully-constructed entry. Updates in-memory state and persists to disk.
53    fn append_entry(&mut self, entry: SessionEntry) -> Result<(), String>;
54
55    /// Look up an entry by ID.
56    fn get_entry(&self, id: &str) -> Option<SessionEntry>;
57
58    /// Find all entries of the given `type` string.
59    fn find_entries(&self, type_name: &str) -> Vec<SessionEntry>;
60
61    /// Get the human-readable label for an entry, if any.
62    fn get_label(&self, id: &str) -> Option<String>;
63
64    /// Get the timestamp of the latest label change for an entry, if any.
65    /// Pi-compatible: used by get_tree() to populate labelTimestamp.
66    fn get_label_timestamp(&self, id: &str) -> Option<String>;
67
68    /// Walk from `leaf_id` (or current leaf, if None) to root, returning entries in path order.
69    fn get_path_to_root(&self, leaf_id: Option<&str>) -> Result<Vec<SessionEntry>, String>;
70
71    /// Return all entries in insertion order.
72    fn get_entries(&self) -> Vec<SessionEntry>;
73
74    /// The file path on disk, if this storage is file-backed.
75    fn path(&self) -> Option<&Path>;
76}
77
78// ── Helpers shared by both implementations ─────────────────────────
79
80/// Given an entry, return the effective leaf ID after it.
81/// For `Leaf` entries, returns `targetId`; for all others, returns `entry.id`.
82fn leaf_id_after_entry(entry: &SessionEntry) -> Option<String> {
83    match entry {
84        SessionEntry::Leaf(e) => e.target_id.clone(),
85        _ => Some(entry.id().to_string()),
86    }
87}
88
89/// Update the label cache from an entry (call after every append).
90/// Pi-compatible: also tracks the timestamp of the latest label change.
91fn update_label_cache(
92    labels_by_id: &mut std::collections::HashMap<String, String>,
93    label_timestamps_by_id: &mut std::collections::HashMap<String, String>,
94    entry: &SessionEntry,
95) {
96    if let SessionEntry::Label(e) = entry {
97        if let Some(label) = &e.label {
98            let trimmed = label.trim();
99            if trimmed.is_empty() {
100                labels_by_id.remove(&e.target_id);
101                label_timestamps_by_id.remove(&e.target_id);
102            } else {
103                labels_by_id.insert(e.target_id.clone(), trimmed.to_string());
104                label_timestamps_by_id.insert(e.target_id.clone(), e.timestamp.clone());
105            }
106        } else {
107            labels_by_id.remove(&e.target_id);
108            label_timestamps_by_id.remove(&e.target_id);
109        }
110    }
111}
112
113/// Build a label cache from a slice of entries.
114fn build_labels_by_id(
115    entries: &[SessionEntry],
116) -> (
117    std::collections::HashMap<String, String>,
118    std::collections::HashMap<String, String>,
119) {
120    let mut labels = std::collections::HashMap::new();
121    let mut timestamps = std::collections::HashMap::new();
122    for entry in entries {
123        update_label_cache(&mut labels, &mut timestamps, entry);
124    }
125    (labels, timestamps)
126}
127
128// ── InMemorySessionStorage ─────────────────────────────────────────
129
130/// Fully in-memory storage — no file I/O.
131/// Pi-compatible: owns all state (entries, labels, leaf).
132pub struct InMemorySessionStorage {
133    metadata: SessionMetadata,
134    entries: Vec<SessionEntry>,
135    by_id: std::collections::HashMap<String, SessionEntry>,
136    labels_by_id: std::collections::HashMap<String, String>,
137    label_timestamps_by_id: std::collections::HashMap<String, String>,
138    leaf_id: Option<String>,
139}
140
141impl InMemorySessionStorage {
142    /// Create empty storage with explicit metadata.
143    pub fn new(metadata: SessionMetadata) -> Self {
144        Self {
145            metadata,
146            entries: Vec::new(),
147            by_id: std::collections::HashMap::new(),
148            labels_by_id: std::collections::HashMap::new(),
149            label_timestamps_by_id: std::collections::HashMap::new(),
150            leaf_id: None,
151        }
152    }
153}
154
155impl SessionStorage for InMemorySessionStorage {
156    fn metadata(&self) -> SessionMetadata {
157        self.metadata.clone()
158    }
159
160    fn get_leaf_id(&self) -> Option<String> {
161        self.leaf_id.clone()
162    }
163
164    fn set_leaf_id(&mut self, leaf_id: Option<&str>) -> Result<(), String> {
165        if let Some(id) = leaf_id
166            && !self.by_id.contains_key(id)
167        {
168            return Err(format!("Entry {} not found", id));
169        }
170        // Pi-compatible: leaf is in-memory only, no LeafEntry created
171        self.leaf_id = leaf_id.map(|s| s.to_string());
172        Ok(())
173    }
174
175    fn create_entry_id(&self) -> String {
176        generate_entry_id(&self.by_id)
177    }
178
179    fn append_entry(&mut self, entry: SessionEntry) -> Result<(), String> {
180        let id = entry.id().to_string();
181        self.by_id.insert(id.clone(), entry);
182        self.entries
183            .push(self.by_id.get(&id).expect("just inserted").clone());
184        self.leaf_id = leaf_id_after_entry(self.by_id.get(&id).expect("just inserted"));
185        update_label_cache(
186            &mut self.labels_by_id,
187            &mut self.label_timestamps_by_id,
188            self.by_id.get(&id).expect("just inserted"),
189        );
190        Ok(())
191    }
192
193    fn get_entry(&self, id: &str) -> Option<SessionEntry> {
194        self.by_id.get(id).cloned()
195    }
196
197    fn find_entries(&self, type_name: &str) -> Vec<SessionEntry> {
198        self.entries
199            .iter()
200            .filter(|e| entry_type_name(e) == type_name)
201            .cloned()
202            .collect()
203    }
204
205    fn get_label(&self, id: &str) -> Option<String> {
206        self.labels_by_id.get(id).cloned()
207    }
208
209    fn get_label_timestamp(&self, id: &str) -> Option<String> {
210        self.label_timestamps_by_id.get(id).cloned()
211    }
212
213    fn get_path_to_root(&self, leaf_id: Option<&str>) -> Result<Vec<SessionEntry>, String> {
214        let start_id = leaf_id.or(self.leaf_id.as_deref());
215        if start_id.is_none() {
216            return Ok(vec![]);
217        }
218        let sid = start_id.unwrap();
219        let mut path: Vec<SessionEntry> = Vec::new();
220        let mut current = self.by_id.get(sid);
221        if current.is_none() {
222            return Err(format!("Entry {} not found", sid));
223        }
224        while let Some(entry) = current {
225            path.push(entry.clone());
226            match entry.parent_id() {
227                Some(pid) => {
228                    current = self.by_id.get(pid);
229                }
230                None => break,
231            }
232        }
233        path.reverse();
234        Ok(path)
235    }
236
237    fn get_entries(&self) -> Vec<SessionEntry> {
238        self.entries.clone()
239    }
240
241    fn path(&self) -> Option<&Path> {
242        None
243    }
244}
245
246// ── JsonlSessionStorage ────────────────────────────────────────────
247
248/// File-backed storage: holds full state in memory and persists to a JSONL file.
249/// Pi-compatible: loads from file on creation, appends on every write.
250pub struct JsonlSessionStorage {
251    metadata: SessionMetadata,
252    file_path: PathBuf,
253    entries: Vec<SessionEntry>,
254    by_id: std::collections::HashMap<String, SessionEntry>,
255    labels_by_id: std::collections::HashMap<String, String>,
256    label_timestamps_by_id: std::collections::HashMap<String, String>,
257    leaf_id: Option<String>,
258}
259
260impl JsonlSessionStorage {
261    /// Create a new session at the given path. Writes the header.
262    pub fn create(
263        file_path: PathBuf,
264        cwd: &str,
265        session_id: &str,
266        parent_session_path: Option<String>,
267    ) -> Result<Self, String> {
268        let created_at = chrono::Utc::now().to_rfc3339();
269        let header = SessionHeader {
270            type_: "session".to_string(),
271            version: Some(crate::agent::session::CURRENT_SESSION_VERSION),
272            id: session_id.to_string(),
273            timestamp: created_at.clone(),
274            cwd: cwd.to_string(),
275            parent_session: parent_session_path.clone(),
276        };
277
278        // Ensure parent directory exists
279        if let Some(parent) = file_path.parent() {
280            std::fs::create_dir_all(parent)
281                .map_err(|e| format!("Failed to create session directory: {}", e))?;
282        }
283
284        // Write header
285        let header_json = serde_json::to_string(&header)
286            .map_err(|e| format!("Failed to serialize header: {}", e))?;
287        std::fs::write(&file_path, header_json + "\n")
288            .map_err(|e| format!("Failed to write session file: {}", e))?;
289
290        let metadata = SessionMetadata {
291            id: session_id.to_string(),
292            created_at,
293            cwd: cwd.to_string(),
294            path: Some(file_path.clone()),
295            parent_session_path,
296        };
297
298        Ok(Self {
299            metadata,
300            file_path,
301            entries: Vec::new(),
302            by_id: std::collections::HashMap::new(),
303            labels_by_id: std::collections::HashMap::new(),
304            label_timestamps_by_id: std::collections::HashMap::new(),
305            leaf_id: None,
306        })
307    }
308
309    /// Open an existing session file. Loads all entries into memory.
310    pub fn open(file_path: PathBuf) -> Result<Self, String> {
311        let (header, entries) = load_session_from_file(&file_path);
312        let header = header
313            .ok_or_else(|| format!("Invalid or missing session header: {}", file_path.display()))?;
314
315        let metadata = SessionMetadata {
316            id: header.id.clone(),
317            created_at: header.timestamp.clone(),
318            cwd: header.cwd,
319            path: Some(file_path.clone()),
320            parent_session_path: header.parent_session,
321        };
322
323        let by_id: std::collections::HashMap<_, _> = entries
324            .iter()
325            .map(|e| (e.id().to_string(), e.clone()))
326            .collect();
327        let (labels_by_id, label_timestamps_by_id) = build_labels_by_id(&entries);
328        let leaf_id = entries.last().and_then(leaf_id_after_entry);
329
330        Ok(Self {
331            metadata,
332            file_path,
333            entries,
334            by_id,
335            labels_by_id,
336            label_timestamps_by_id,
337            leaf_id,
338        })
339    }
340
341    /// Append a line to the file.
342    fn append_to_file(&self, entry: &SessionEntry) -> Result<(), String> {
343        append_entry_to_file(&self.file_path, entry)
344            .map_err(|e| format!("Failed to append session entry: {}", e))
345    }
346}
347
348impl SessionStorage for JsonlSessionStorage {
349    fn metadata(&self) -> SessionMetadata {
350        self.metadata.clone()
351    }
352
353    fn get_leaf_id(&self) -> Option<String> {
354        self.leaf_id.clone()
355    }
356
357    fn set_leaf_id(&mut self, leaf_id: Option<&str>) -> Result<(), String> {
358        if let Some(id) = leaf_id
359            && !self.by_id.contains_key(id)
360        {
361            return Err(format!("Entry {} not found", id));
362        }
363        // Pi-compatible: leaf is in-memory only, no LeafEntry written to file
364        self.leaf_id = leaf_id.map(|s| s.to_string());
365        Ok(())
366    }
367
368    fn create_entry_id(&self) -> String {
369        generate_entry_id(&self.by_id)
370    }
371
372    fn append_entry(&mut self, entry: SessionEntry) -> Result<(), String> {
373        self.append_to_file(&entry)?;
374        let id = entry.id().to_string();
375        self.by_id.insert(id.clone(), entry);
376        self.entries
377            .push(self.by_id.get(&id).expect("just inserted").clone());
378        self.leaf_id = leaf_id_after_entry(self.by_id.get(&id).expect("just inserted"));
379        update_label_cache(
380            &mut self.labels_by_id,
381            &mut self.label_timestamps_by_id,
382            self.by_id.get(&id).expect("just inserted"),
383        );
384        Ok(())
385    }
386
387    fn get_entry(&self, id: &str) -> Option<SessionEntry> {
388        self.by_id.get(id).cloned()
389    }
390
391    fn find_entries(&self, type_name: &str) -> Vec<SessionEntry> {
392        self.entries
393            .iter()
394            .filter(|e| entry_type_name(e) == type_name)
395            .cloned()
396            .collect()
397    }
398
399    fn get_label(&self, id: &str) -> Option<String> {
400        self.labels_by_id.get(id).cloned()
401    }
402
403    fn get_label_timestamp(&self, id: &str) -> Option<String> {
404        self.label_timestamps_by_id.get(id).cloned()
405    }
406
407    fn get_path_to_root(&self, leaf_id: Option<&str>) -> Result<Vec<SessionEntry>, String> {
408        let start_id = leaf_id.or(self.leaf_id.as_deref());
409        if start_id.is_none() {
410            return Ok(vec![]);
411        }
412        let sid = start_id.unwrap();
413        let mut path: Vec<SessionEntry> = Vec::new();
414        let mut current = self.by_id.get(sid);
415        if current.is_none() {
416            return Err(format!("Entry {} not found", sid));
417        }
418        while let Some(entry) = current {
419            path.push(entry.clone());
420            match entry.parent_id() {
421                Some(pid) => {
422                    current = self.by_id.get(pid);
423                }
424                None => break,
425            }
426        }
427        path.reverse();
428        Ok(path)
429    }
430
431    fn get_entries(&self) -> Vec<SessionEntry> {
432        self.entries.clone()
433    }
434
435    fn path(&self) -> Option<&Path> {
436        Some(&self.file_path)
437    }
438}
439
440// ── Helper: entry type name ────────────────────────────────────────
441
442/// Return the type string for a SessionEntry (pi-compatible).
443fn entry_type_name(entry: &SessionEntry) -> &'static str {
444    match entry {
445        SessionEntry::Message(_) => "message",
446        SessionEntry::ThinkingLevelChange(_) => "thinking_level_change",
447        SessionEntry::ModelChange(_) => "model_change",
448        SessionEntry::ActiveToolsChange(_) => "active_tools_change",
449        SessionEntry::Compaction(_) => "compaction",
450        SessionEntry::BranchSummary(_) => "branch_summary",
451        SessionEntry::SessionInfo(_) => "session_info",
452        SessionEntry::Label(_) => "label",
453        SessionEntry::Custom(_) => "custom",
454        SessionEntry::CustomMessage(_) => "custom_message",
455        SessionEntry::Leaf(_) => "leaf",
456    }
457}
458
459#[cfg(test)]
460mod tests {
461    use super::super::model::{MessageCost, MessageEntry};
462    use super::*;
463    use crate::agent::types::user_message;
464    use tempfile::TempDir;
465
466    fn make_session_meta(id: &str) -> SessionMetadata {
467        SessionMetadata {
468            id: id.to_string(),
469            created_at: chrono::Utc::now().to_rfc3339(),
470            cwd: "/tmp/test".to_string(),
471            path: None,
472            parent_session_path: None,
473        }
474    }
475
476    fn make_msg_entry(id: &str, parent: Option<&str>, text: &str) -> SessionEntry {
477        SessionEntry::Message(MessageEntry {
478            id: id.to_string(),
479            parent_id: parent.map(|s| s.to_string()),
480            timestamp: chrono::Utc::now().to_rfc3339(),
481            message: user_message(text),
482            cost: MessageCost::ZERO,
483        })
484    }
485
486    // ── InMemorySessionStorage tests ──────────────────────────────────
487
488    #[test]
489    fn test_in_memory_empty() {
490        let meta = make_session_meta("test");
491        let storage = InMemorySessionStorage::new(meta.clone());
492        assert_eq!(storage.metadata().id, "test");
493        assert!(storage.get_leaf_id().is_none());
494        assert!(storage.get_entries().is_empty());
495    }
496
497    #[test]
498    fn test_in_memory_append_and_get() {
499        let mut storage = InMemorySessionStorage::new(make_session_meta("s1"));
500        let e = make_msg_entry("m1", None, "hello");
501        storage.append_entry(e).unwrap();
502        assert_eq!(storage.get_leaf_id(), Some("m1".to_string()));
503        assert_eq!(storage.get_entry("m1").unwrap().id(), "m1");
504        assert_eq!(storage.get_entries().len(), 1);
505    }
506
507    #[test]
508    fn test_in_memory_path_to_root() {
509        let mut storage = InMemorySessionStorage::new(make_session_meta("s1"));
510        storage
511            .append_entry(make_msg_entry("m1", None, "first"))
512            .unwrap();
513        storage
514            .append_entry(make_msg_entry("m2", Some("m1"), "second"))
515            .unwrap();
516        storage
517            .append_entry(make_msg_entry("m3", Some("m2"), "third"))
518            .unwrap();
519
520        let path = storage.get_path_to_root(Some("m3")).unwrap();
521        assert_eq!(path.len(), 3);
522        assert_eq!(path[0].id(), "m1");
523        assert_eq!(path[2].id(), "m3");
524    }
525
526    #[test]
527    fn test_in_memory_labels() {
528        let mut storage = InMemorySessionStorage::new(make_session_meta("s1"));
529        storage
530            .append_entry(make_msg_entry("m1", None, "first"))
531            .unwrap();
532
533        // Add label
534        let label_entry = SessionEntry::Label(crate::agent::session::LabelEntry {
535            id: "l1".to_string(),
536            parent_id: Some("m1".to_string()),
537            timestamp: chrono::Utc::now().to_rfc3339(),
538            target_id: "m1".to_string(),
539            label: Some("important".to_string()),
540        });
541        storage.append_entry(label_entry).unwrap();
542        assert_eq!(storage.get_label("m1"), Some("important".to_string()));
543
544        // Remove label
545        let unlabel_entry = SessionEntry::Label(crate::agent::session::LabelEntry {
546            id: "l2".to_string(),
547            parent_id: Some("l1".to_string()),
548            timestamp: chrono::Utc::now().to_rfc3339(),
549            target_id: "m1".to_string(),
550            label: None,
551        });
552        storage.append_entry(unlabel_entry).unwrap();
553        assert_eq!(storage.get_label("m1"), None);
554    }
555
556    #[test]
557    fn test_in_memory_set_leaf_id() {
558        let mut storage = InMemorySessionStorage::new(make_session_meta("s1"));
559        storage
560            .append_entry(make_msg_entry("m1", None, "first"))
561            .unwrap();
562        storage
563            .append_entry(make_msg_entry("m2", Some("m1"), "second"))
564            .unwrap();
565
566        // Set leaf to m1 (branching) — in-memory only, no LeafEntry written
567        storage.set_leaf_id(Some("m1")).unwrap();
568        assert_eq!(storage.get_leaf_id(), Some("m1".to_string()));
569
570        // Verify no LeafEntry was added to entries
571        let entries = storage.get_entries();
572        assert_eq!(entries.len(), 2);
573        assert!(!entries.iter().any(|e| matches!(e, SessionEntry::Leaf(_))));
574    }
575
576    #[test]
577    fn test_in_memory_find_entries() {
578        let mut storage = InMemorySessionStorage::new(make_session_meta("s1"));
579        storage
580            .append_entry(make_msg_entry("m1", None, "first"))
581            .unwrap();
582        let tl =
583            SessionEntry::ThinkingLevelChange(crate::agent::session::ThinkingLevelChangeEntry {
584                id: "tc1".to_string(),
585                parent_id: Some("m1".to_string()),
586                timestamp: chrono::Utc::now().to_rfc3339(),
587                thinking_level: "high".to_string(),
588            });
589        storage.append_entry(tl).unwrap();
590        storage
591            .append_entry(make_msg_entry("m2", Some("tc1"), "second"))
592            .unwrap();
593
594        let msgs = storage.find_entries("message");
595        assert_eq!(msgs.len(), 2);
596        let tls = storage.find_entries("thinking_level_change");
597        assert_eq!(tls.len(), 1);
598    }
599
600    // ── JsonlSessionStorage tests ────────────────────────────────────
601
602    #[test]
603    fn test_jsonl_create_and_append() {
604        let tmp = TempDir::new().unwrap();
605        let path = tmp.path().join("session.jsonl");
606
607        let mut storage =
608            JsonlSessionStorage::create(path.clone(), "/tmp/test", "s1", None).unwrap();
609        assert_eq!(storage.metadata().id, "s1");
610        assert!(path.exists());
611
612        storage
613            .append_entry(make_msg_entry("m1", None, "hello"))
614            .unwrap();
615        assert_eq!(storage.get_entries().len(), 1);
616        assert_eq!(storage.get_leaf_id(), Some("m1".to_string()));
617
618        // Verify persistence by opening again
619        let loaded = JsonlSessionStorage::open(path).unwrap();
620        assert_eq!(loaded.get_entries().len(), 1);
621        assert_eq!(loaded.get_entry("m1").unwrap().id(), "m1");
622    }
623
624    #[test]
625    fn test_jsonl_open_and_traverse() {
626        let tmp = TempDir::new().unwrap();
627        let path = tmp.path().join("session.jsonl");
628
629        let mut storage =
630            JsonlSessionStorage::create(path.clone(), "/tmp/test", "s1", None).unwrap();
631        storage
632            .append_entry(make_msg_entry("m1", None, "first"))
633            .unwrap();
634        storage
635            .append_entry(make_msg_entry("m2", Some("m1"), "second"))
636            .unwrap();
637        drop(storage);
638
639        let loaded = JsonlSessionStorage::open(path).unwrap();
640        let path_to = loaded.get_path_to_root(Some("m2")).unwrap();
641        assert_eq!(path_to.len(), 2);
642    }
643
644    #[test]
645    fn test_in_memory_label_timestamp() {
646        let mut storage = InMemorySessionStorage::new(make_session_meta("s1"));
647        storage
648            .append_entry(make_msg_entry("m1", None, "first"))
649            .unwrap();
650
651        // No label yet
652        assert!(storage.get_label_timestamp("m1").is_none());
653
654        // Add label
655        let label_entry = SessionEntry::Label(crate::agent::session::LabelEntry {
656            id: "l1".to_string(),
657            parent_id: Some("m1".to_string()),
658            timestamp: "2026-06-30T12:00:00Z".to_string(),
659            target_id: "m1".to_string(),
660            label: Some("star".to_string()),
661        });
662        storage.append_entry(label_entry).unwrap();
663        assert_eq!(storage.get_label("m1"), Some("star".to_string()));
664        assert_eq!(
665            storage.get_label_timestamp("m1").as_deref(),
666            Some("2026-06-30T12:00:00Z")
667        );
668
669        // Remove label
670        let unlabel_entry = SessionEntry::Label(crate::agent::session::LabelEntry {
671            id: "l2".to_string(),
672            parent_id: Some("l1".to_string()),
673            timestamp: "2026-06-30T13:00:00Z".to_string(),
674            target_id: "m1".to_string(),
675            label: None,
676        });
677        storage.append_entry(unlabel_entry).unwrap();
678        assert!(storage.get_label_timestamp("m1").is_none());
679    }
680
681    #[test]
682    fn test_jsonl_label_timestamp_persistence() {
683        let tmp = TempDir::new().unwrap();
684        let path = tmp.path().join("session.jsonl");
685
686        let mut storage =
687            JsonlSessionStorage::create(path.clone(), "/tmp/test", "s1", None).unwrap();
688        storage
689            .append_entry(make_msg_entry("m1", None, "first"))
690            .unwrap();
691
692        let label_entry = SessionEntry::Label(crate::agent::session::LabelEntry {
693            id: "l1".to_string(),
694            parent_id: Some("m1".to_string()),
695            timestamp: "2026-06-30T12:00:00Z".to_string(),
696            target_id: "m1".to_string(),
697            label: Some("important".to_string()),
698        });
699        storage.append_entry(label_entry).unwrap();
700        drop(storage);
701
702        // Reopen and verify timestamp survived
703        let loaded = JsonlSessionStorage::open(path).unwrap();
704        assert_eq!(loaded.get_label("m1"), Some("important".to_string()));
705        assert_eq!(
706            loaded.get_label_timestamp("m1").as_deref(),
707            Some("2026-06-30T12:00:00Z")
708        );
709    }
710}
rab/agent/session/storage.rs

rab/agent/session/
storage.rs
