Skip to main content

baml_agent/session/
mod.rs

1//! Session persistence — Claude Code compatible JSONL format.
2//!
3//! Modules:
4//! - `traits` — MessageRole, AgentMessage, EntryType
5//! - `format` — PersistedMessage, parse_entry, make_persisted
6//! - `time` — ISO timestamps, UUID v7 extraction, truncation
7//! - `store` — Session struct (CRUD, trimming)
8//! - `meta` — SessionMeta, list/search/import
9
10pub(crate) mod format;
11mod meta;
12mod store;
13pub(crate) mod time;
14pub mod traits;
15
16#[cfg(feature = "search")]
17pub use meta::search_sessions;
18pub use meta::{import_claude_session, list_sessions, SessionMeta};
19pub use store::Session;
20pub use traits::{AgentMessage, EntryType, MessageRole};
21
22#[cfg(test)]
23pub(crate) mod tests {
24    use super::format::{make_persisted, parse_entry};
25    use super::time::{now_iso, truncate_str, truncate_topic, uuid_v7_timestamp};
26    use super::*;
27
28    #[derive(Clone, Debug, PartialEq)]
29    pub(crate) enum TestRole {
30        System,
31        User,
32        Assistant,
33        Tool,
34    }
35
36    impl MessageRole for TestRole {
37        fn system() -> Self {
38            Self::System
39        }
40        fn user() -> Self {
41            Self::User
42        }
43        fn assistant() -> Self {
44            Self::Assistant
45        }
46        fn tool() -> Self {
47            Self::Tool
48        }
49        fn as_str(&self) -> &str {
50            match self {
51                Self::System => "system",
52                Self::User => "user",
53                Self::Assistant => "assistant",
54                Self::Tool => "tool",
55            }
56        }
57        fn parse_role(s: &str) -> Option<Self> {
58            match s {
59                "system" => Some(Self::System),
60                "user" => Some(Self::User),
61                "assistant" => Some(Self::Assistant),
62                "tool" => Some(Self::Tool),
63                _ => None,
64            }
65        }
66    }
67
68    #[derive(Clone)]
69    pub(crate) struct TestMsg {
70        pub role: TestRole,
71        pub content: String,
72    }
73
74    impl AgentMessage for TestMsg {
75        type Role = TestRole;
76        fn new(role: TestRole, content: String) -> Self {
77            Self { role, content }
78        }
79        fn role(&self) -> &TestRole {
80            &self.role
81        }
82        fn content(&self) -> &str {
83            &self.content
84        }
85    }
86
87    // --- EntryType ---
88
89    #[test]
90    fn entry_type_roundtrip() {
91        for t in [
92            EntryType::User,
93            EntryType::Assistant,
94            EntryType::System,
95            EntryType::Tool,
96        ] {
97            let json = serde_json::to_string(&t).unwrap();
98            let back: EntryType = serde_json::from_str(&json).unwrap();
99            assert_eq!(t, back);
100        }
101    }
102
103    #[test]
104    fn entry_type_rejects_invalid() {
105        assert!(EntryType::parse("progress").is_none());
106        assert!(EntryType::parse("file-history-snapshot").is_none());
107        assert!(EntryType::parse("").is_none());
108    }
109
110    // --- Format ---
111
112    #[test]
113    fn user_message_serialized_as_plain_string() {
114        let p = make_persisted(EntryType::User, "hello", "sid", None);
115        let json: serde_json::Value = serde_json::to_value(&p).unwrap();
116        assert!(json["message"]["content"].is_string());
117        assert_eq!(json["message"]["content"].as_str(), Some("hello"));
118        assert_eq!(json["message"]["role"].as_str(), Some("user"));
119    }
120
121    #[test]
122    fn assistant_message_serialized_as_blocks() {
123        let p = make_persisted(EntryType::Assistant, "thinking...", "sid", None);
124        let json: serde_json::Value = serde_json::to_value(&p).unwrap();
125        let blocks = json["message"]["content"].as_array().unwrap();
126        assert_eq!(blocks.len(), 1);
127        assert_eq!(blocks[0]["type"].as_str(), Some("text"));
128        assert_eq!(blocks[0]["text"].as_str(), Some("thinking..."));
129    }
130
131    #[test]
132    fn system_message_serialized_as_plain_string() {
133        let p = make_persisted(EntryType::System, "you are an agent", "sid", None);
134        let json: serde_json::Value = serde_json::to_value(&p).unwrap();
135        assert!(json["message"]["content"].is_string());
136    }
137
138    // --- parse_entry ---
139
140    #[test]
141    fn parse_entry_claude_code_user_format() {
142        let entry: serde_json::Value = serde_json::from_str(
143            r#"{"type":"user","message":{"role":"user","content":"fix the bug"},"uuid":"abc","sessionId":"s1","timestamp":"2026-03-07T10:00:00.000Z"}"#
144        ).unwrap();
145        let (et, content) = parse_entry(&entry).unwrap();
146        assert_eq!(et, EntryType::User);
147        assert_eq!(content, "fix the bug");
148    }
149
150    #[test]
151    fn parse_entry_claude_code_assistant_format() {
152        let entry: serde_json::Value = serde_json::from_str(
153            r#"{"type":"assistant","message":{"role":"assistant","content":[{"type":"text","text":"Looking at it..."},{"type":"tool_use","name":"Read","id":"x","input":{}}]},"uuid":"def","sessionId":"s1","timestamp":"2026-03-07T10:00:01.000Z"}"#
154        ).unwrap();
155        let (et, content) = parse_entry(&entry).unwrap();
156        assert_eq!(et, EntryType::Assistant);
157        assert!(content.contains("Looking at it..."));
158        assert!(content.contains("[tool: Read]"));
159    }
160
161    #[test]
162    fn parse_entry_skips_thinking_only() {
163        let entry: serde_json::Value = serde_json::from_str(
164            r#"{"type":"assistant","message":{"role":"assistant","content":[{"type":"thinking","thinking":"hmm..."}]},"uuid":"ghi","sessionId":"s1","timestamp":"2026-03-07T10:00:02.000Z"}"#
165        ).unwrap();
166        assert!(parse_entry(&entry).is_none());
167    }
168
169    #[test]
170    fn parse_entry_skips_progress() {
171        let entry: serde_json::Value = serde_json::from_str(
172            r#"{"type":"progress","data":{"type":"hook_progress"},"uuid":"a"}"#,
173        )
174        .unwrap();
175        assert!(parse_entry(&entry).is_none());
176    }
177
178    #[test]
179    fn parse_entry_legacy_format() {
180        let entry: serde_json::Value =
181            serde_json::from_str(r#"{"role":"user","content":"hello legacy"}"#).unwrap();
182        let (et, content) = parse_entry(&entry).unwrap();
183        assert_eq!(et, EntryType::User);
184        assert_eq!(content, "hello legacy");
185    }
186
187    // --- Time & truncation ---
188
189    #[test]
190    fn now_iso_produces_valid_timestamp() {
191        let ts = now_iso();
192        assert!(ts.ends_with('Z'));
193        assert_eq!(ts.len(), 24);
194        assert_eq!(&ts[4..5], "-");
195        assert_eq!(&ts[10..11], "T");
196    }
197
198    #[test]
199    fn truncate_str_ascii() {
200        assert_eq!(truncate_str("hello world", 5), "hello");
201        assert_eq!(truncate_str("hi", 10), "hi");
202    }
203
204    #[test]
205    fn truncate_str_utf8_safe() {
206        let s = "ab\u{00e9}cd\u{00fc}ef";
207        let t = truncate_str(s, 4);
208        assert!(t.len() <= 4);
209        assert_eq!(t, "ab\u{00e9}");
210
211        let t2 = truncate_str(s, 3);
212        assert!(t2.len() <= 3);
213        assert_eq!(t2, "ab");
214    }
215
216    #[test]
217    fn truncate_str_emoji() {
218        let s = "Hello \u{1f30d}\u{1f30d}\u{1f30d}";
219        let t = truncate_str(s, 8);
220        assert!(t.len() <= 8);
221        assert_eq!(t, "Hello ");
222    }
223
224    #[test]
225    fn truncate_topic_multibyte() {
226        let long_multibyte = "\u{00e9}".repeat(200);
227        let topic = truncate_topic(&long_multibyte);
228        assert!(topic.ends_with("..."));
229        assert!(topic.len() <= 120);
230    }
231
232    // --- UUID v7 ---
233
234    #[test]
235    fn uuid_v7_is_time_ordered() {
236        let id1 = uuid::Uuid::now_v7().to_string();
237        std::thread::sleep(std::time::Duration::from_millis(2));
238        let id2 = uuid::Uuid::now_v7().to_string();
239        assert!(
240            id2 > id1,
241            "v7 UUIDs should be time-ordered: {} > {}",
242            id2,
243            id1
244        );
245    }
246
247    #[test]
248    fn uuid_v7_timestamp_extraction() {
249        let before = std::time::SystemTime::now()
250            .duration_since(std::time::UNIX_EPOCH)
251            .unwrap()
252            .as_secs();
253        let id = uuid::Uuid::now_v7().to_string();
254        let after = std::time::SystemTime::now()
255            .duration_since(std::time::UNIX_EPOCH)
256            .unwrap()
257            .as_secs();
258        let ts = uuid_v7_timestamp(&id).unwrap();
259        assert!(ts >= before && ts <= after);
260    }
261
262    #[test]
263    fn uuid_v7_timestamp_invalid() {
264        assert!(uuid_v7_timestamp("not-a-uuid").is_none());
265        assert!(uuid_v7_timestamp("550e8400-e29b-41d4-a716-446655440000").is_none());
266    }
267
268    // --- Session CRUD ---
269
270    #[test]
271    fn trim_preserves_system_and_recent() {
272        let dir = std::env::temp_dir().join("baml_mod_test_trim");
273        let _ = std::fs::remove_dir_all(&dir);
274        let mut session = Session::<TestMsg>::new(dir.to_str().unwrap(), 10).unwrap();
275
276        session.push(TestRole::System, "sys prompt".into());
277        for i in 0..20 {
278            let role = if i % 2 == 0 {
279                TestRole::User
280            } else {
281                TestRole::Assistant
282            };
283            session.push(role, format!("msg {}", i));
284        }
285        assert_eq!(session.len(), 21);
286
287        let trimmed = session.trim();
288        assert!(trimmed > 0);
289        assert!(session.len() <= 12);
290        assert_eq!(session.messages()[0].role(), &TestRole::System);
291        assert!(session.messages()[1].content().contains("trimmed"));
292        assert_eq!(session.messages().last().unwrap().content(), "msg 19");
293
294        let _ = std::fs::remove_dir_all(&dir);
295    }
296
297    #[test]
298    fn trim_noop_small_history() {
299        let dir = std::env::temp_dir().join("baml_mod_test_noop");
300        let _ = std::fs::remove_dir_all(&dir);
301        let mut session = Session::<TestMsg>::new(dir.to_str().unwrap(), 60).unwrap();
302        session.push(TestRole::User, "hello".into());
303        assert_eq!(session.trim(), 0);
304        let _ = std::fs::remove_dir_all(&dir);
305    }
306
307    #[test]
308    fn persist_and_reload() {
309        let dir = std::env::temp_dir().join("baml_mod_test_persist");
310        let _ = std::fs::remove_dir_all(&dir);
311        let mut session = Session::<TestMsg>::new(dir.to_str().unwrap(), 60).unwrap();
312        let sid = session.session_id().to_string();
313        session.push(TestRole::User, "hello world".into());
314        session.push(TestRole::Assistant, "hi there".into());
315
316        let path = session.session_file().to_path_buf();
317        let loaded = Session::<TestMsg>::resume(&path, dir.to_str().unwrap(), 60);
318        assert_eq!(loaded.len(), 2);
319        assert_eq!(loaded.messages()[0].content(), "hello world");
320        assert_eq!(loaded.messages()[1].role(), &TestRole::Assistant);
321        assert_eq!(loaded.session_id(), sid);
322
323        // Verify user = plain string, assistant = blocks
324        let raw = std::fs::read_to_string(&path).unwrap();
325        let first: serde_json::Value = serde_json::from_str(raw.lines().next().unwrap()).unwrap();
326        assert!(first["message"]["content"].is_string());
327        let second: serde_json::Value = serde_json::from_str(raw.lines().nth(1).unwrap()).unwrap();
328        assert!(second["message"]["content"].is_array());
329
330        let _ = std::fs::remove_dir_all(&dir);
331    }
332
333    #[test]
334    fn persist_parent_uuid_chain() {
335        let dir = std::env::temp_dir().join("baml_mod_test_parent");
336        let _ = std::fs::remove_dir_all(&dir);
337        let mut session = Session::<TestMsg>::new(dir.to_str().unwrap(), 60).unwrap();
338        session.push(TestRole::User, "first".into());
339        session.push(TestRole::Assistant, "second".into());
340        session.push(TestRole::User, "third".into());
341
342        let raw = std::fs::read_to_string(session.session_file()).unwrap();
343        let entries: Vec<serde_json::Value> = raw
344            .lines()
345            .filter_map(|l| serde_json::from_str(l).ok())
346            .collect();
347
348        assert!(entries[0]["parentUuid"].is_null());
349        assert_eq!(
350            entries[1]["parentUuid"].as_str(),
351            entries[0]["uuid"].as_str()
352        );
353        assert_eq!(
354            entries[2]["parentUuid"].as_str(),
355            entries[1]["uuid"].as_str()
356        );
357
358        let _ = std::fs::remove_dir_all(&dir);
359    }
360
361    #[test]
362    fn persist_multibyte_content() {
363        let dir = std::env::temp_dir().join("baml_mod_test_multibyte");
364        let _ = std::fs::remove_dir_all(&dir);
365        let mut session = Session::<TestMsg>::new(dir.to_str().unwrap(), 60).unwrap();
366        session.push(
367            TestRole::User,
368            "caf\u{00e9} na\u{00ef}ve r\u{00e9}sum\u{00e9}".into(),
369        );
370        session.push(TestRole::Assistant, "got it! \u{1f389}".into());
371
372        let path = session.session_file().to_path_buf();
373        let loaded = Session::<TestMsg>::resume(&path, dir.to_str().unwrap(), 60);
374        assert_eq!(
375            loaded.messages()[0].content(),
376            "caf\u{00e9} na\u{00ef}ve r\u{00e9}sum\u{00e9}"
377        );
378        assert_eq!(loaded.messages()[1].content(), "got it! \u{1f389}");
379
380        let _ = std::fs::remove_dir_all(&dir);
381    }
382
383    #[test]
384    fn load_legacy_format() {
385        let dir = std::env::temp_dir().join("baml_mod_test_legacy");
386        let _ = std::fs::remove_dir_all(&dir);
387        std::fs::create_dir_all(&dir).unwrap();
388
389        let path = dir.join("session_1234567890.jsonl");
390        let legacy = [
391            r#"{"role":"user","content":"hello legacy"}"#,
392            r#"{"role":"assistant","content":"hi from old format"}"#,
393        ];
394        std::fs::write(&path, legacy.join("\n")).unwrap();
395
396        let loaded = Session::<TestMsg>::resume(&path, dir.to_str().unwrap(), 60);
397        assert_eq!(loaded.len(), 2);
398        assert_eq!(loaded.messages()[0].content(), "hello legacy");
399        assert_eq!(loaded.session_id(), "session_1234567890");
400
401        let _ = std::fs::remove_dir_all(&dir);
402    }
403
404    #[test]
405    fn resume_last_finds_latest() {
406        let dir = std::env::temp_dir().join("baml_mod_test_resume");
407        let _ = std::fs::remove_dir_all(&dir);
408
409        let mut s1 = Session::<TestMsg>::new(dir.to_str().unwrap(), 60).unwrap();
410        s1.push(TestRole::User, "first".into());
411        std::thread::sleep(std::time::Duration::from_millis(10));
412        let mut s2 = Session::<TestMsg>::new(dir.to_str().unwrap(), 60).unwrap();
413        s2.push(TestRole::User, "second".into());
414
415        let resumed = Session::<TestMsg>::resume_last(dir.to_str().unwrap(), 60).unwrap();
416        assert_eq!(resumed.messages()[0].content(), "second");
417
418        let _ = std::fs::remove_dir_all(&dir);
419    }
420
421    // --- SessionMeta ---
422
423    #[test]
424    fn session_meta_extracts_topic() {
425        let dir = std::env::temp_dir().join("baml_mod_test_topic");
426        let _ = std::fs::remove_dir_all(&dir);
427
428        let mut s = Session::<TestMsg>::new(dir.to_str().unwrap(), 60).unwrap();
429        s.push(TestRole::System, "you are an agent".into());
430        s.push(TestRole::User, "deploy to production".into());
431
432        let meta = SessionMeta::from_path(s.session_file()).unwrap();
433        assert_eq!(meta.topic, "deploy to production");
434        assert_eq!(meta.message_count, 2);
435        assert!(meta.size_bytes > 0);
436        assert!(meta.session_id.is_some());
437
438        let _ = std::fs::remove_dir_all(&dir);
439    }
440
441    #[test]
442    fn session_meta_created_from_uuid_v7() {
443        let dir = std::env::temp_dir().join("baml_mod_test_uuid_ts");
444        let _ = std::fs::remove_dir_all(&dir);
445
446        let before = std::time::SystemTime::now()
447            .duration_since(std::time::UNIX_EPOCH)
448            .unwrap()
449            .as_secs();
450        let mut s = Session::<TestMsg>::new(dir.to_str().unwrap(), 60).unwrap();
451        s.push(TestRole::User, "test".into());
452        let after = std::time::SystemTime::now()
453            .duration_since(std::time::UNIX_EPOCH)
454            .unwrap()
455            .as_secs();
456
457        let meta = SessionMeta::from_path(s.session_file()).unwrap();
458        assert!(meta.created >= before && meta.created <= after);
459
460        let _ = std::fs::remove_dir_all(&dir);
461    }
462
463    #[test]
464    fn list_sessions_returns_sorted() {
465        let dir = std::env::temp_dir().join("baml_mod_test_list");
466        let _ = std::fs::remove_dir_all(&dir);
467
468        let mut s1 = Session::<TestMsg>::new(dir.to_str().unwrap(), 60).unwrap();
469        s1.push(TestRole::User, "fix parser bug".into());
470        std::thread::sleep(std::time::Duration::from_millis(10));
471        let mut s2 = Session::<TestMsg>::new(dir.to_str().unwrap(), 60).unwrap();
472        s2.push(TestRole::User, "add new feature".into());
473
474        let sessions = list_sessions(dir.to_str().unwrap());
475        assert_eq!(sessions.len(), 2);
476        assert!(sessions[0].created >= sessions[1].created);
477
478        let _ = std::fs::remove_dir_all(&dir);
479    }
480
481    #[test]
482    fn list_sessions_empty_dir() {
483        let dir = std::env::temp_dir().join("baml_mod_test_empty");
484        let _ = std::fs::remove_dir_all(&dir);
485        let _ = std::fs::create_dir_all(&dir);
486        assert!(list_sessions(dir.to_str().unwrap()).is_empty());
487        let _ = std::fs::remove_dir_all(&dir);
488    }
489
490    #[test]
491    fn topic_truncated_for_long_messages() {
492        let dir = std::env::temp_dir().join("baml_mod_test_long_topic");
493        let _ = std::fs::remove_dir_all(&dir);
494
495        let mut s = Session::<TestMsg>::new(dir.to_str().unwrap(), 60).unwrap();
496        s.push(TestRole::User, "a".repeat(200));
497
498        let meta = SessionMeta::from_path(s.session_file()).unwrap();
499        assert!(meta.topic.len() <= 120);
500        assert!(meta.topic.ends_with("..."));
501
502        let _ = std::fs::remove_dir_all(&dir);
503    }
504
505    // --- Search ---
506
507    #[cfg(feature = "search")]
508    #[test]
509    fn search_sessions_fuzzy() {
510        let dir = std::env::temp_dir().join("baml_mod_test_search");
511        let _ = std::fs::remove_dir_all(&dir);
512
513        let mut s1 = Session::<TestMsg>::new(dir.to_str().unwrap(), 60).unwrap();
514        s1.push(TestRole::User, "fix parser bug in baml".into());
515        std::thread::sleep(std::time::Duration::from_millis(10));
516        let mut s2 = Session::<TestMsg>::new(dir.to_str().unwrap(), 60).unwrap();
517        s2.push(TestRole::User, "deploy to production".into());
518        std::thread::sleep(std::time::Duration::from_millis(10));
519        let mut s3 = Session::<TestMsg>::new(dir.to_str().unwrap(), 60).unwrap();
520        s3.push(TestRole::User, "fix loop detection bug".into());
521
522        let results = search_sessions(dir.to_str().unwrap(), "fix bug");
523        assert!(!results.is_empty());
524        let topics: Vec<&str> = results.iter().map(|(_, m)| m.topic.as_str()).collect();
525        assert!(topics.iter().any(|t| t.contains("fix")));
526
527        let _ = std::fs::remove_dir_all(&dir);
528    }
529
530    // --- Import ---
531
532    #[test]
533    fn import_claude_session_converts() {
534        let dir = std::env::temp_dir().join("baml_mod_test_import");
535        let _ = std::fs::remove_dir_all(&dir);
536        std::fs::create_dir_all(&dir).unwrap();
537
538        let claude_session = dir.join("claude_session.jsonl");
539        let lines = [
540            r#"{"type":"progress","data":{"type":"hook_progress"},"uuid":"a","timestamp":"2026-03-07"}"#,
541            r#"{"type":"user","message":{"role":"user","content":"fix the parser bug"},"uuid":"b","sessionId":"test-sid","timestamp":"2026-03-07"}"#,
542            r#"{"type":"assistant","message":{"role":"assistant","content":[{"type":"text","text":"Looking at the parser..."},{"type":"tool_use","name":"Read","id":"x","input":{}}]},"uuid":"c","sessionId":"test-sid","timestamp":"2026-03-07"}"#,
543            r#"{"type":"system","message":{"content":"context info"},"uuid":"d","sessionId":"test-sid","timestamp":"2026-03-07"}"#,
544        ];
545        std::fs::write(&claude_session, lines.join("\n")).unwrap();
546
547        let out_dir = dir.join("sessions");
548        let result = import_claude_session(&claude_session, out_dir.to_str().unwrap());
549        assert!(result.is_some());
550
551        let output = result.unwrap();
552        let content = std::fs::read_to_string(&output).unwrap();
553        let entries: Vec<serde_json::Value> = content
554            .lines()
555            .filter_map(|l| serde_json::from_str(l).ok())
556            .collect();
557
558        assert_eq!(entries.len(), 3);
559        assert_eq!(entries[0]["type"].as_str(), Some("user"));
560        assert_eq!(entries[1]["type"].as_str(), Some("assistant"));
561        assert_eq!(entries[2]["type"].as_str(), Some("system"));
562
563        let sid = entries[0]["sessionId"].as_str().unwrap();
564        assert_ne!(sid, "test-sid");
565
566        let _ = std::fs::remove_dir_all(&dir);
567    }
568
569    // --- Integration ---
570
571    #[test]
572    fn load_real_claude_session() {
573        let Ok(home) = std::env::var("HOME") else {
574            return;
575        };
576        let claude_dir = std::path::Path::new(&home).join(".claude/projects");
577        if !claude_dir.exists() {
578            return;
579        }
580
581        let Some(project_dir) = std::fs::read_dir(&claude_dir).ok().and_then(|rd| {
582            rd.filter_map(|e| e.ok()).find(|e| {
583                e.path().is_dir()
584                    && std::fs::read_dir(e.path())
585                        .ok()
586                        .map(|rd2| {
587                            rd2.filter_map(|e2| e2.ok())
588                                .any(|e2| e2.path().extension().is_some_and(|ext| ext == "jsonl"))
589                        })
590                        .unwrap_or(false)
591            })
592        }) else {
593            return;
594        };
595
596        let smallest = std::fs::read_dir(project_dir.path())
597            .unwrap()
598            .filter_map(|e| e.ok())
599            .filter(|e| e.path().extension().is_some_and(|ext| ext == "jsonl"))
600            .min_by_key(|e| e.metadata().map(|m| m.len()).unwrap_or(u64::MAX));
601        let Some(smallest) = smallest else { return };
602
603        let out_dir = std::env::temp_dir().join("baml_mod_test_real");
604        let _ = std::fs::remove_dir_all(&out_dir);
605
606        let result = import_claude_session(&smallest.path(), out_dir.to_str().unwrap());
607        if let Some(output) = result {
608            let content = std::fs::read_to_string(&output).unwrap();
609            assert!(content.lines().count() > 0);
610            for line in content.lines() {
611                let v: serde_json::Value = serde_json::from_str(line).unwrap();
612                assert!(v["type"].as_str().is_some());
613                assert!(v["sessionId"].as_str().is_some());
614            }
615        }
616
617        let _ = std::fs::remove_dir_all(&out_dir);
618    }
619}