1pub(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 #[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 #[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 #[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 #[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 #[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 #[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 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 #[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 #[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 #[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 #[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}