1use std::fs::File;
2use std::path::{Path, PathBuf};
3use std::sync::LazyLock;
4
5use agent_base::{AgentError, AgentResult};
6use fs2::FileExt;
7use regex::Regex;
8
9static SESSION_ID_RE: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"^[a-zA-Z0-9\-_]+$").unwrap());
10static SNAPSHOT_NAME_RE: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"^[a-zA-Z0-9\-_]+$").unwrap());
11
12#[derive(Debug)]
16pub struct SessionContext {
17 pub session_id: String,
19 pub session_dir: PathBuf,
21 pub base_dir: PathBuf,
23 pub is_new_session: bool,
25 _lock: Option<File>,
26}
27
28impl SessionContext {
29 #[allow(dead_code)]
31 pub fn session_id_path(&self) -> PathBuf {
32 self.session_dir.join("session_id")
33 }
34
35 pub fn metadata_path(&self) -> PathBuf {
37 self.session_dir.join("session_meta.json")
38 }
39
40 pub fn log_path(&self) -> PathBuf {
42 self.session_dir.join("session.log")
43 }
44
45 pub fn turn_path(&self, turn: usize) -> PathBuf {
47 self.session_dir.join(format!("turn_{:03}.jsonl", turn))
48 }
49
50 pub fn last_turn_number(&self) -> u32 {
58 let Ok(rd) = std::fs::read_dir(&self.session_dir) else {
59 return 0;
60 };
61 let mut max = 0u32;
62 for entry in rd.flatten() {
63 let name = entry.file_name();
64 let name = name.to_string_lossy();
65 if let Some(n) =
66 name.strip_prefix("turn_").and_then(|s| s.strip_suffix(".jsonl")).and_then(|s| s.parse::<u32>().ok())
67 {
68 max = max.max(n);
69 }
70 }
71 max
72 }
73}
74
75pub fn validate_session_id(session_id: &str) -> AgentResult<()> {
79 if session_id.is_empty() || session_id.len() > 128 {
80 return Err(AgentError::config_error(format!("Session ID must be 1-128 characters, got {}", session_id.len())));
81 }
82
83 if !SESSION_ID_RE.is_match(session_id) {
84 return Err(AgentError::config_error(format!(
85 "Invalid session_id format '{}'. Only alphanumeric, hyphens, and underscores allowed.",
86 session_id
87 )));
88 }
89
90 Ok(())
91}
92
93pub fn validate_snapshot_name(name: &str) -> AgentResult<()> {
98 if name.is_empty() || name.len() > 64 {
99 return Err(AgentError::config_error(format!("Snapshot name must be 1-64 characters, got {}", name.len())));
100 }
101
102 if name.contains('/') || name.contains('\\') || name.contains("..") {
103 return Err(AgentError::config_error(format!(
104 "Invalid snapshot name '{}': path separators and '..' are not allowed",
105 name
106 )));
107 }
108
109 if !SNAPSHOT_NAME_RE.is_match(name) {
110 return Err(AgentError::config_error(format!(
111 "Invalid snapshot name '{}'. Only alphanumeric, hyphens, and underscores allowed.",
112 name
113 )));
114 }
115
116 Ok(())
117}
118
119pub fn resolve_session_id(cli_session_id: Option<&str>) -> AgentResult<String> {
121 if let Some(id) = cli_session_id {
122 validate_session_id(id)?;
123 return Ok(id.to_string());
124 }
125
126 if let Ok(id) = std::env::var("PHI_SESSION_ID")
127 && !id.is_empty()
128 {
129 validate_session_id(&id)?;
130 return Ok(id);
131 }
132
133 Ok(generate_session_id())
134}
135
136pub fn generate_session_id() -> String {
138 let now = chrono::Local::now();
139 let uuid = uuid::Uuid::new_v4().to_string();
140 let uuid_short = &uuid[..8.min(uuid.len())];
141 format!("{}_{}", now.format("%Y%m%d"), uuid_short)
142}
143
144pub fn get_or_create_session_dir(session_id: &str, base_dir: &Path) -> AgentResult<(PathBuf, bool)> {
148 let session_dir = base_dir.join("sessions").join(session_id);
149 let is_new = !session_dir.exists();
150
151 if is_new {
152 std::fs::create_dir_all(&session_dir)?;
153 tracing::info!(session_id = %session_id, path = %session_dir.display(), "created new session directory");
154 } else {
155 tracing::info!(session_id = %session_id, path = %session_dir.display(), "reusing existing session directory");
156 }
157
158 std::fs::write(session_dir.join("session_id"), session_id)?;
160
161 update_session_meta(&session_dir, session_id)?;
163
164 Ok((session_dir, is_new))
165}
166
167pub fn acquire_session_lock(session_dir: &Path) -> AgentResult<File> {
172 let lock_path = session_dir.join("session.lock");
173 let file = File::create(&lock_path)?;
174
175 file.try_lock_exclusive().map_err(|_| {
176 AgentError::resource_unavailable(format!(
177 "Session '{}' is currently in use by another process",
178 session_dir.file_name().map(|n| n.to_string_lossy().to_string()).unwrap_or_default()
179 ))
180 })?;
181
182 Ok(file)
183}
184
185fn update_session_meta(session_dir: &Path, session_id: &str) -> AgentResult<()> {
187 let meta_path = session_dir.join("session_meta.json");
188
189 let mut meta = if meta_path.exists() {
190 let content = std::fs::read_to_string(&meta_path)?;
191 serde_json::from_str::<serde_json::Value>(&content)?
192 } else {
193 serde_json::json!({
194 "session_id": session_id,
195 "created_at": chrono::Utc::now().to_rfc3339(),
196 })
197 };
198
199 meta["last_active_at"] = serde_json::json!(chrono::Utc::now().to_rfc3339());
200
201 std::fs::write(&meta_path, serde_json::to_string_pretty(&meta)?)?;
202 Ok(())
203}
204
205pub fn cleanup_expired_sessions(base_dir: &Path, max_age_days: i64) -> AgentResult<u32> {
210 let sessions_dir = base_dir.join("sessions");
211 if !sessions_dir.exists() {
212 return Ok(0);
213 }
214
215 let now = chrono::Utc::now();
216 let mut cleaned = 0;
217
218 for entry in std::fs::read_dir(&sessions_dir)? {
219 let entry = entry?;
220 let path = entry.path();
221
222 if !path.is_dir() {
223 continue;
224 }
225
226 let lock_path = path.join("session.lock");
227 if lock_path.exists()
228 && let Ok(file) = File::open(&lock_path)
229 && file.try_lock_shared().is_err()
230 {
231 continue; }
233
234 let meta_path = path.join("session_meta.json");
235 if !meta_path.exists() {
236 std::fs::remove_dir_all(&path)?;
237 cleaned += 1;
238 continue;
239 }
240
241 let content = std::fs::read_to_string(&meta_path)?;
242 let meta: serde_json::Value = serde_json::from_str(&content)?;
243
244 if let Some(last_active) = meta["last_active_at"].as_str()
245 && let Ok(last_active) = chrono::DateTime::parse_from_rfc3339(last_active)
246 {
247 let age = now - last_active.with_timezone(&chrono::Utc);
248 if age.num_days() > max_age_days {
249 tracing::info!(path = %path.display(), age_days = age.num_days(), "removing expired session");
250 std::fs::remove_dir_all(&path)?;
251 cleaned += 1;
252 }
253 }
254 }
255
256 if cleaned > 0 {
257 tracing::info!(count = cleaned, "cleaned up expired sessions");
258 }
259
260 Ok(cleaned)
261}
262
263pub fn create_snapshot(session_ctx: &SessionContext, name: &str, base_dir: &Path) -> AgentResult<PathBuf> {
270 validate_snapshot_name(name)?;
271 let snapshot_dir = base_dir.join("snapshots").join(name);
272
273 if snapshot_dir.exists() {
274 std::fs::remove_dir_all(&snapshot_dir)?;
275 }
276 std::fs::create_dir_all(&snapshot_dir)?;
277
278 let meta_path = session_ctx.metadata_path();
280 if meta_path.exists() {
281 std::fs::copy(&meta_path, snapshot_dir.join("session_meta.json"))?;
282 }
283
284 let mut turn_count = 0u32;
286 for entry in std::fs::read_dir(&session_ctx.session_dir)? {
287 let entry = entry?;
288 let name = entry.file_name();
289 let name_str = name.to_string_lossy();
290 if name_str.starts_with("turn_") && name_str.ends_with(".jsonl") {
291 std::fs::copy(entry.path(), snapshot_dir.join(&*name_str))?;
292 turn_count += 1;
293 }
294 }
295
296 let info = serde_json::json!({
298 "session_id": session_ctx.session_id,
299 "snapshot_name": name,
300 "created_at": chrono::Utc::now().to_rfc3339(),
301 "turn_count": turn_count,
302 });
303 std::fs::write(snapshot_dir.join("snapshot_info.json"), serde_json::to_string_pretty(&info)?)?;
304
305 tracing::info!(
306 session_id = %session_ctx.session_id,
307 name = name,
308 turns = turn_count,
309 "session snapshot created"
310 );
311
312 Ok(snapshot_dir)
313}
314
315pub fn list_snapshots(base_dir: &Path) -> AgentResult<Vec<SnapshotInfo>> {
319 let snapshots_dir = base_dir.join("snapshots");
320 if !snapshots_dir.exists() {
321 return Ok(Vec::new());
322 }
323
324 let mut snapshots = Vec::new();
325 for entry in std::fs::read_dir(&snapshots_dir)? {
326 let entry = entry?;
327 let path = entry.path();
328 if !path.is_dir() {
329 continue;
330 }
331
332 let info_path = path.join("snapshot_info.json");
333 if !info_path.exists() {
334 continue;
335 }
336
337 let content = std::fs::read_to_string(&info_path)?;
338 let info: serde_json::Value = serde_json::from_str(&content)?;
339
340 let name = path.file_name().map(|n| n.to_string_lossy().to_string()).unwrap_or_default();
341 let session_id = info["session_id"].as_str().unwrap_or("-").to_string();
342 let created_at = info["created_at"].as_str().unwrap_or("-").to_string();
343 let turn_count = info["turn_count"].as_u64().unwrap_or(0) as u32;
344
345 snapshots.push(SnapshotInfo { name, session_id, created_at, turn_count, path });
346 }
347
348 snapshots.sort_by(|a, b| b.created_at.cmp(&a.created_at));
349 Ok(snapshots)
350}
351
352pub struct SnapshotInfo {
354 pub name: String,
356 pub session_id: String,
358 pub created_at: String,
360 pub turn_count: u32,
362 pub path: PathBuf,
364}
365
366pub fn restore_snapshot(name: &str, base_dir: &Path) -> AgentResult<SessionContext> {
370 validate_snapshot_name(name)?;
371 let snapshot_dir = base_dir.join("snapshots").join(name);
372 if !snapshot_dir.exists() {
373 return Err(AgentError::config_error(format!("Snapshot '{}' not found at {}", name, snapshot_dir.display())));
374 }
375
376 let new_session_id = generate_session_id();
378 let (new_session_dir, _) = get_or_create_session_dir(&new_session_id, base_dir)?;
379
380 for entry in std::fs::read_dir(&snapshot_dir)? {
382 let entry = entry?;
383 let fname = entry.file_name();
384 let name_str = fname.to_string_lossy();
385 if name_str.starts_with("turn_") && name_str.ends_with(".jsonl") {
386 std::fs::copy(entry.path(), new_session_dir.join(&*name_str))?;
387 }
388 }
389
390 let lock = acquire_session_lock(&new_session_dir)?;
391 tracing::info!(from = name, to = %new_session_id, "session restored from snapshot");
392
393 Ok(SessionContext {
394 session_id: new_session_id,
395 session_dir: new_session_dir,
396 base_dir: base_dir.to_path_buf(),
397 is_new_session: false,
398 _lock: Some(lock),
399 })
400}
401
402pub fn delete_snapshot(name: &str, base_dir: &Path) -> AgentResult<()> {
404 validate_snapshot_name(name)?;
405 let snapshot_dir = base_dir.join("snapshots").join(name);
406 if !snapshot_dir.exists() {
407 return Err(AgentError::config_error(format!("Snapshot '{}' not found", name)));
408 }
409 std::fs::remove_dir_all(&snapshot_dir)?;
410 tracing::info!(name = name, "snapshot deleted");
411 Ok(())
412}
413
414pub fn resolve_session(cli_session_id: Option<&str>, base_dir: &Path) -> AgentResult<SessionContext> {
419 let session_id = resolve_session_id(cli_session_id)?;
420 let (session_dir, is_new) = get_or_create_session_dir(&session_id, base_dir)?;
421 let lock = acquire_session_lock(&session_dir)?;
422
423 Ok(SessionContext {
424 session_id,
425 session_dir,
426 base_dir: base_dir.to_path_buf(),
427 is_new_session: is_new,
428 _lock: Some(lock),
429 })
430}
431
432#[cfg(test)]
433mod tests {
434 use super::*;
435 use tempfile::TempDir;
436
437 #[test]
438 fn test_validate_session_id_valid() {
439 assert!(validate_session_id("my-session-123").is_ok());
440 assert!(validate_session_id("test_456").is_ok());
441 assert!(validate_session_id("a").is_ok());
442 }
443
444 #[test]
445 fn test_validate_session_id_invalid() {
446 assert!(validate_session_id("").is_err());
447 assert!(validate_session_id("my session").is_err());
448 assert!(validate_session_id("../etc").is_err());
449 assert!(validate_session_id("path/traversal").is_err());
450 }
451
452 #[test]
453 fn test_generate_session_id() {
454 let id = generate_session_id();
455 assert!(id.contains('_'));
456 let parts: Vec<&str> = id.split('_').collect();
457 assert_eq!(parts.len(), 2);
458 assert_eq!(parts[0].len(), 8);
459 assert_eq!(parts[1].len(), 8);
460 }
461
462 #[test]
463 fn test_session_context_methods() {
464 let tmp = TempDir::new().unwrap();
465 let ctx = resolve_session(Some("test-ctx"), tmp.path()).unwrap();
466
467 assert_eq!(ctx.session_id, "test-ctx");
468 assert!(ctx.session_id_path().exists());
469 assert!(ctx.metadata_path().exists());
470 assert_eq!(ctx.log_path(), ctx.session_dir.join("session.log"));
471 assert_eq!(ctx.turn_path(1), ctx.session_dir.join("turn_001.jsonl"));
472 }
473
474 #[test]
475 fn test_last_turn_number_scans_existing_turns() {
476 let tmp = TempDir::new().unwrap();
477 let ctx = resolve_session(Some("turn-scan"), tmp.path()).unwrap();
478 assert_eq!(ctx.last_turn_number(), 0, "empty session should start at 0");
479
480 std::fs::write(ctx.turn_path(1), "").unwrap();
481 std::fs::write(ctx.turn_path(3), "").unwrap();
482 assert_eq!(ctx.last_turn_number(), 3);
483
484 std::fs::write(ctx.session_dir.join("session.log"), "").unwrap();
486 std::fs::write(ctx.session_dir.join("turn_abc.jsonl"), "").unwrap();
487 std::fs::write(ctx.session_dir.join("turn_004.jsonl"), "").unwrap();
488 assert_eq!(ctx.last_turn_number(), 4);
489 }
490
491 #[test]
492 fn test_cleanup_expired_sessions() {
493 let tmp = TempDir::new().unwrap();
494 let (dir, _) = get_or_create_session_dir("old-session", tmp.path()).unwrap();
495
496 let meta_path = dir.join("session_meta.json");
497 let mut meta: serde_json::Value = serde_json::from_str(&std::fs::read_to_string(&meta_path).unwrap()).unwrap();
498 let old = (chrono::Utc::now() - chrono::Duration::days(8)).to_rfc3339();
499 meta["last_active_at"] = serde_json::json!(old);
500 std::fs::write(&meta_path, serde_json::to_string_pretty(&meta).unwrap()).unwrap();
501
502 get_or_create_session_dir("new-session", tmp.path()).unwrap();
503 let cleaned = cleanup_expired_sessions(tmp.path(), 7).unwrap();
504 assert_eq!(cleaned, 1);
505 assert!(!dir.exists());
506 }
507
508 #[test]
511 fn test_validate_session_id_returns_config_error() {
512 let err = validate_session_id("").unwrap_err();
514 assert!(matches!(err, AgentError::ConfigError(_)), "expected ConfigError for empty session ID, got {:?}", err);
515
516 let err = validate_session_id("../etc").unwrap_err();
518 assert!(
519 matches!(err, AgentError::ConfigError(_)),
520 "expected ConfigError for invalid session ID, got {:?}",
521 err
522 );
523 }
524
525 #[test]
526 fn test_resolve_session_returns_config_error_for_invalid_id() {
527 let tmp = TempDir::new().unwrap();
528 let result = resolve_session(Some("bad id with spaces"), tmp.path());
529 assert!(result.is_err());
530 let err = result.err().unwrap();
531 assert!(
532 matches!(err, AgentError::ConfigError(_)),
533 "expected ConfigError for invalid session ID, got {:?}",
534 err
535 );
536 }
537
538 #[test]
539 fn test_acquire_session_lock_returns_resource_unavailable_when_locked() {
540 let tmp = TempDir::new().unwrap();
541 let (_dir, _) = get_or_create_session_dir("lock-test", tmp.path()).unwrap();
542 let session_dir = tmp.path().join("sessions").join("lock-test");
543
544 let _lock1 = acquire_session_lock(&session_dir).unwrap();
546
547 let err = acquire_session_lock(&session_dir).unwrap_err();
549 assert!(
550 matches!(err, AgentError::ResourceUnavailable(_)),
551 "expected ResourceUnavailable for locked session, got {:?}",
552 err
553 );
554 }
555
556 #[test]
559 fn test_validate_snapshot_name_valid() {
560 assert!(validate_snapshot_name("my-snapshot").is_ok());
561 assert!(validate_snapshot_name("test_123").is_ok());
562 assert!(validate_snapshot_name("a").is_ok());
563 assert!(validate_snapshot_name("snapshot-2024-01-01").is_ok());
564 }
565
566 #[test]
567 fn test_validate_snapshot_name_invalid() {
568 assert!(validate_snapshot_name("").is_err());
569 assert!(validate_snapshot_name("../etc").is_err());
570 assert!(validate_snapshot_name("path/traversal").is_err());
571 assert!(validate_snapshot_name("back\\slash").is_err());
572 assert!(validate_snapshot_name("dot..dot").is_err());
573 assert!(validate_snapshot_name("has space").is_err());
574 assert!(validate_snapshot_name(&"x".repeat(65)).is_err()); }
576
577 #[test]
578 fn test_validate_snapshot_name_returns_config_error() {
579 let err = validate_snapshot_name("").unwrap_err();
580 assert!(matches!(err, AgentError::ConfigError(_)), "expected ConfigError, got {:?}", err);
581
582 let err = validate_snapshot_name("../etc").unwrap_err();
583 assert!(matches!(err, AgentError::ConfigError(_)), "expected ConfigError for path traversal, got {:?}", err);
584 }
585
586 #[test]
589 fn test_create_snapshot_and_list() {
590 let tmp = TempDir::new().unwrap();
591 let ctx = resolve_session(Some("snap-session"), tmp.path()).unwrap();
592
593 let turn_path = ctx.turn_path(1);
595 std::fs::write(&turn_path, "{}").unwrap();
596
597 let snap_path = create_snapshot(&ctx, "test-snap", tmp.path()).unwrap();
599 assert!(snap_path.exists());
600 assert!(snap_path.join("snapshot_info.json").exists());
601 assert!(snap_path.join("turn_001.jsonl").exists());
602
603 let snaps = list_snapshots(tmp.path()).unwrap();
605 assert_eq!(snaps.len(), 1);
606 assert_eq!(snaps[0].name, "test-snap");
607 assert_eq!(snaps[0].session_id, "snap-session");
608 assert_eq!(snaps[0].turn_count, 1);
609 }
610
611 #[test]
612 fn test_create_snapshot_overwrites_existing() {
613 let tmp = TempDir::new().unwrap();
614 let ctx = resolve_session(Some("snap-session"), tmp.path()).unwrap();
615
616 create_snapshot(&ctx, "dup-snap", tmp.path()).unwrap();
617 let result = create_snapshot(&ctx, "dup-snap", tmp.path());
619 assert!(result.is_ok());
620
621 let snaps = list_snapshots(tmp.path()).unwrap();
622 assert_eq!(snaps.len(), 1);
623 }
624
625 #[test]
626 fn test_list_snapshots_empty_dir() {
627 let tmp = TempDir::new().unwrap();
628 let snaps = list_snapshots(tmp.path()).unwrap();
629 assert!(snaps.is_empty());
630 }
631
632 #[test]
633 fn test_list_snapshots_sorted_by_date_desc() {
634 let tmp = TempDir::new().unwrap();
635 let ctx = resolve_session(Some("snap-session"), tmp.path()).unwrap();
636
637 create_snapshot(&ctx, "first", tmp.path()).unwrap();
638 std::thread::sleep(std::time::Duration::from_millis(10));
639 create_snapshot(&ctx, "second", tmp.path()).unwrap();
640
641 let snaps = list_snapshots(tmp.path()).unwrap();
642 assert_eq!(snaps.len(), 2);
643 assert_eq!(snaps[0].name, "second");
645 assert_eq!(snaps[1].name, "first");
646 }
647
648 #[test]
651 fn test_restore_snapshot_success() {
652 let tmp = TempDir::new().unwrap();
653 let ctx = resolve_session(Some("snap-session"), tmp.path()).unwrap();
654
655 std::fs::write(ctx.turn_path(1), r#"{"type":"text_delta","text":"hello"}"#).unwrap();
657 create_snapshot(&ctx, "restore-me", tmp.path()).unwrap();
658
659 let restored = restore_snapshot("restore-me", tmp.path()).unwrap();
661 assert_ne!(restored.session_id, "snap-session");
662 assert!(!restored.is_new_session);
663 assert!(restored.turn_path(1).exists());
664
665 let content = std::fs::read_to_string(restored.turn_path(1)).unwrap();
667 assert!(content.contains("hello"));
668 }
669
670 #[test]
671 fn test_restore_snapshot_not_found() {
672 let tmp = TempDir::new().unwrap();
673 let err = restore_snapshot("no-such-snapshot", tmp.path()).unwrap_err();
674 assert!(matches!(err, AgentError::ConfigError(_)), "expected ConfigError for missing snapshot, got {:?}", err);
675 }
676
677 #[test]
680 fn test_delete_snapshot_success() {
681 let tmp = TempDir::new().unwrap();
682 let ctx = resolve_session(Some("snap-session"), tmp.path()).unwrap();
683
684 create_snapshot(&ctx, "del-me", tmp.path()).unwrap();
685 assert_eq!(list_snapshots(tmp.path()).unwrap().len(), 1);
686
687 delete_snapshot("del-me", tmp.path()).unwrap();
688 assert_eq!(list_snapshots(tmp.path()).unwrap().len(), 0);
689 }
690
691 #[test]
692 fn test_delete_snapshot_not_found() {
693 let tmp = TempDir::new().unwrap();
694 let err = delete_snapshot("no-such-snapshot", tmp.path()).unwrap_err();
695 assert!(matches!(err, AgentError::ConfigError(_)), "expected ConfigError for missing snapshot, got {:?}", err);
696 }
697
698 #[test]
701 fn test_session_context_stores_base_dir() {
702 let tmp = TempDir::new().unwrap();
703 let ctx = resolve_session(Some("base-test"), tmp.path()).unwrap();
704 assert_eq!(ctx.base_dir, tmp.path());
705 assert!(ctx.session_dir.starts_with(&ctx.base_dir));
706 }
707}
708
709#[cfg(test)]
710mod proptests {
711 use super::*;
712
713 proptest::proptest! {
714 #[test]
715 fn validate_session_id_never_panics(s in ".*") {
716 let _ = validate_session_id(&s);
717 }
718
719 #[test]
720 fn validate_session_id_accepts_valid_chars_only(s in "[a-zA-Z0-9_-]{1,128}") {
721 proptest::prop_assert!(validate_session_id(&s).is_ok(), "valid session_id '{}' should pass", s);
722 }
723
724 #[test]
725 fn validate_snapshot_name_rejects_path_traversal(s in ".*") {
726 let result = validate_snapshot_name(&s);
727 if s.contains('/') || s.contains('\\') || s.contains("..") {
728 proptest::prop_assert!(result.is_err(), "snapshot name '{}' with path traversal should fail", s);
729 }
730 }
731
732 #[test]
733 fn cleanup_timestamp_rfc3339_never_panics(ts in "[0-9T:+ -Zz.]{1,40}") {
734 let _ = chrono::DateTime::parse_from_rfc3339(&ts);
735 }
736 }
737}