1use std::collections::BTreeSet;
2use std::fs;
3use std::path::{Component, Path, PathBuf};
4use std::time::{Duration, SystemTime, UNIX_EPOCH};
5
6use anyhow::{Context, Result};
7use base64::Engine as _;
8use base64::engine::general_purpose::STANDARD as BASE64;
9use serde::{Deserialize, Serialize};
10
11use crate::utils::error_messages::ERR_CREATE_CHECKPOINT_DIR;
12use crate::utils::file_utils::{ensure_dir_exists, ensure_dir_exists_sync, write_json_file};
13use crate::utils::path::canonicalize_workspace;
14use crate::utils::session_archive::SessionMessage;
15
16const MAX_DESCRIPTION_LEN: usize = 160;
17use crate::core::SECONDS_PER_DAY;
18pub const DEFAULT_CHECKPOINTS_ENABLED: bool = true;
19pub const DEFAULT_MAX_SNAPSHOTS: usize = 50;
20pub const DEFAULT_MAX_AGE_DAYS: u64 = 30;
21
22fn normalized_prompt_text(text: &str) -> Option<&str> {
23 let trimmed = text.trim();
24 (!trimmed.is_empty()).then_some(trimmed)
25}
26
27fn sanitize_relative_path(path: &Path) -> Option<PathBuf> {
28 if path.is_absolute() {
29 return None;
30 }
31
32 let mut normalized = PathBuf::new();
33 for component in path.components() {
34 match component {
35 Component::CurDir => {}
36 Component::Normal(part) => normalized.push(part),
37 Component::ParentDir => {
38 if !normalized.pop() {
39 return None;
40 }
41 }
42 Component::Prefix(_) | Component::RootDir => {
43 return None;
44 }
45 }
46 }
47 Some(normalized)
48}
49
50#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
51pub struct SnapshotMetadata {
52 pub id: String,
53 pub turn_number: usize,
54 pub created_at: u64,
55 pub description: String,
56 pub message_count: usize,
57 pub file_count: usize,
58 #[serde(default, skip_serializing_if = "Option::is_none")]
59 pub prompt_text: Option<String>,
60 #[serde(default, skip_serializing_if = "Option::is_none")]
61 pub prompt_message_index: Option<usize>,
62}
63
64impl SnapshotMetadata {
65 pub fn resolved_prompt_text<'a>(
66 &'a self,
67 conversation: &'a [SessionMessage],
68 ) -> Option<String> {
69 self.prompt_text
70 .as_deref()
71 .and_then(normalized_prompt_text)
72 .map(str::to_string)
73 .or_else(|| SnapshotManager::derive_prompt_metadata(conversation).0)
74 }
75}
76
77#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
78pub enum FileEncoding {
79 Utf8,
80 Base64,
81}
82
83#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
84pub struct FileSnapshot {
85 pub path: String,
86 pub deleted: bool,
87 #[serde(skip_serializing_if = "Option::is_none")]
88 pub encoding: Option<FileEncoding>,
89 #[serde(skip_serializing_if = "Option::is_none")]
90 pub data: Option<String>,
91}
92
93#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
94pub struct StoredSnapshot {
95 pub metadata: SnapshotMetadata,
96 pub conversation: Vec<SessionMessage>,
97 pub files: Vec<FileSnapshot>,
98}
99
100#[derive(Debug, Clone, Copy, PartialEq, Eq)]
101pub enum RevertScope {
102 Conversation,
103 Code,
104 Both,
105}
106
107impl RevertScope {
108 pub fn includes_code(self) -> bool {
109 matches!(self, Self::Code | Self::Both)
110 }
111
112 pub fn includes_conversation(self) -> bool {
113 matches!(self, Self::Conversation | Self::Both)
114 }
115}
116
117pub struct SnapshotConfig {
118 pub enabled: bool,
119 pub workspace: PathBuf,
120 pub storage_dir: Option<PathBuf>,
121 pub max_snapshots: usize,
122 pub max_age_days: Option<u64>,
123}
124
125impl SnapshotConfig {
126 pub fn new(workspace: PathBuf) -> Self {
127 Self {
128 enabled: DEFAULT_CHECKPOINTS_ENABLED,
129 workspace,
130 storage_dir: None,
131 max_snapshots: DEFAULT_MAX_SNAPSHOTS,
132 max_age_days: Some(DEFAULT_MAX_AGE_DAYS),
133 }
134 }
135
136 fn storage_dir(&self) -> PathBuf {
137 self.storage_dir
138 .clone()
139 .unwrap_or_else(|| self.workspace.join(".vtcode").join("checkpoints"))
140 }
141}
142
143pub struct SnapshotManager {
144 enabled: bool,
145 workspace: PathBuf,
146 canonical_workspace: PathBuf,
147 storage_dir: PathBuf,
148 max_snapshots: usize,
149 max_age_days: Option<u64>,
150}
151
152impl SnapshotManager {
153 pub fn new(config: SnapshotConfig) -> Result<Self> {
154 let storage_dir = config.storage_dir();
155 let canonical_workspace = canonicalize_workspace(&config.workspace);
156
157 if config.enabled {
158 ensure_dir_exists_sync(&storage_dir).with_context(|| {
159 format!("{}: {}", ERR_CREATE_CHECKPOINT_DIR, storage_dir.display())
160 })?;
161 }
162 Ok(Self {
163 enabled: config.enabled,
164 workspace: config.workspace,
165 canonical_workspace,
166 storage_dir,
167 max_snapshots: config.max_snapshots,
168 max_age_days: config.max_age_days,
169 })
170 }
171
172 pub fn enabled(&self) -> bool {
173 self.enabled
174 }
175
176 fn snapshot_path(&self, turn_number: usize) -> PathBuf {
177 self.storage_dir.join(format!("turn_{}.json", turn_number))
178 }
179
180 fn normalize_path(&self, path: &Path) -> Option<PathBuf> {
181 if path.is_absolute() {
182 if let Ok(canonical_path) = fs::canonicalize(path)
183 && let Ok(stripped) = canonical_path.strip_prefix(&self.canonical_workspace)
184 {
185 return sanitize_relative_path(stripped);
186 }
187
188 if let Ok(stripped) = path.strip_prefix(&self.workspace) {
189 return sanitize_relative_path(stripped);
190 }
191
192 None
193 } else {
194 sanitize_relative_path(path)
195 }
196 }
197
198 fn read_snapshot_files(&self) -> Result<Vec<(usize, PathBuf)>> {
199 let mut entries = Vec::with_capacity(64); if !self.storage_dir.exists() {
201 return Ok(entries);
202 }
203 for entry in fs::read_dir(&self.storage_dir).with_context(|| {
204 format!(
205 "failed to read checkpoint directory: {}",
206 self.storage_dir.display()
207 )
208 })? {
209 let entry = entry?;
210 let path = entry.path();
211 if path.extension().and_then(|ext| ext.to_str()) != Some("json") {
212 continue;
213 }
214 let stem = match path.file_stem().and_then(|stem| stem.to_str()) {
215 Some(value) => value,
216 None => continue,
217 };
218 let turn_str = match stem.strip_prefix("turn_") {
219 Some(value) => value,
220 None => continue,
221 };
222 if let Ok(turn) = turn_str.parse::<usize>() {
223 entries.push((turn, path));
224 }
225 }
226 entries.sort_by_key(|(turn, _)| *turn);
227 Ok(entries)
228 }
229
230 fn encode_file(bytes: &[u8]) -> (FileEncoding, String) {
231 match std::str::from_utf8(bytes) {
232 Ok(text) => (FileEncoding::Utf8, text.to_string()),
233 Err(_) => (FileEncoding::Base64, BASE64.encode(bytes)),
234 }
235 }
236
237 fn decode_file(encoding: FileEncoding, data: &str) -> Result<Vec<u8>> {
238 match encoding {
239 FileEncoding::Utf8 => Ok(data.as_bytes().to_vec()),
240 FileEncoding::Base64 => BASE64
241 .decode(data)
242 .context("failed to decode base64 file contents"),
243 }
244 }
245
246 fn truncate_description(description: &str) -> String {
247 let first_line = description.lines().next().unwrap_or("").trim();
248 vtcode_commons::formatting::truncate_within(first_line, MAX_DESCRIPTION_LEN, "…")
249 }
250
251 fn derive_prompt_metadata(conversation: &[SessionMessage]) -> (Option<String>, Option<usize>) {
252 conversation
253 .iter()
254 .enumerate()
255 .rev()
256 .find_map(|(index, message)| {
257 if message.role != crate::llm::provider::MessageRole::User {
258 return None;
259 }
260
261 let prompt = message.content.as_text();
262 normalized_prompt_text(prompt.as_ref())
263 .map(|prompt| (Some(prompt.to_string()), Some(index)))
264 })
265 .unwrap_or((None, None))
266 }
267
268 fn resolve_prompt_metadata(
269 prompt_text: Option<&str>,
270 prompt_message_index: Option<usize>,
271 conversation: &[SessionMessage],
272 ) -> (Option<String>, Option<usize>) {
273 let (derived_prompt_text, derived_prompt_index) =
274 Self::derive_prompt_metadata(conversation);
275 let prompt_text = prompt_text
276 .and_then(normalized_prompt_text)
277 .map(str::to_string)
278 .or(derived_prompt_text);
279 let prompt_message_index = prompt_message_index
280 .filter(|index| *index < conversation.len())
281 .or(derived_prompt_index);
282 (prompt_text, prompt_message_index)
283 }
284
285 fn hydrate_prompt_metadata(stored: &mut StoredSnapshot) {
286 let (prompt_text, prompt_message_index) = Self::resolve_prompt_metadata(
287 stored.metadata.prompt_text.as_deref(),
288 stored.metadata.prompt_message_index,
289 &stored.conversation,
290 );
291 stored.metadata.prompt_text = prompt_text;
292 stored.metadata.prompt_message_index = prompt_message_index;
293 }
294
295 fn current_timestamp() -> Result<u64> {
296 Ok(SystemTime::now()
297 .duration_since(UNIX_EPOCH)
298 .context("system clock before UNIX_EPOCH")?
299 .as_secs())
300 }
301
302 pub fn next_turn_number(&self) -> Result<usize> {
303 Ok(self
304 .read_snapshot_files()?
305 .into_iter()
306 .map(|(turn, _)| turn)
307 .max()
308 .unwrap_or(0)
309 .saturating_add(1))
310 }
311
312 pub async fn create_snapshot(
313 &self,
314 turn_number: usize,
315 description: &str,
316 conversation: &[SessionMessage],
317 modified_files: &BTreeSet<PathBuf>,
318 prompt_text: Option<&str>,
319 prompt_message_index: Option<usize>,
320 ) -> Result<Option<SnapshotMetadata>> {
321 if !self.enabled {
322 return Ok(None);
323 }
324
325 let timestamp = Self::current_timestamp()?;
326 let mut files = Vec::with_capacity(modified_files.len()); for path in modified_files {
329 let relative = match self.normalize_path(path) {
330 Some(value) => value,
331 None => continue,
332 };
333 let absolute = self.workspace.join(&relative);
334 if tokio::fs::try_exists(&absolute).await.unwrap_or(false) {
335 let bytes = tokio::fs::read(&absolute).await.with_context(|| {
336 format!("failed to read file for checkpoint: {}", absolute.display())
337 })?;
338 let (encoding, data) = Self::encode_file(&bytes);
339 files.push(FileSnapshot {
340 path: relative.to_string_lossy().replace('\\', "/"),
341 deleted: false,
342 encoding: Some(encoding),
343 data: Some(data),
344 });
345 } else {
346 files.push(FileSnapshot {
347 path: relative.to_string_lossy().replace('\\', "/"),
348 deleted: true,
349 encoding: None,
350 data: None,
351 });
352 }
353 }
354
355 let (prompt_text, prompt_message_index) =
356 Self::resolve_prompt_metadata(prompt_text, prompt_message_index, conversation);
357 let description_source = prompt_text.as_deref().unwrap_or(description);
358 let metadata = SnapshotMetadata {
359 id: format!("turn_{}", turn_number),
360 turn_number,
361 created_at: timestamp,
362 description: Self::truncate_description(description_source),
363 message_count: conversation.len(),
364 file_count: files.len(),
365 prompt_text,
366 prompt_message_index,
367 };
368
369 let stored = StoredSnapshot {
370 metadata: metadata.clone(),
371 conversation: conversation.to_vec(),
372 files,
373 };
374
375 let path = self.snapshot_path(turn_number);
376 if let Some(parent) = path.parent() {
377 ensure_dir_exists(parent).await.with_context(|| {
378 format!(
379 "failed to ensure checkpoint directory: {}",
380 parent.display()
381 )
382 })?;
383 }
384
385 write_json_file(&path, &stored)
386 .await
387 .with_context(|| format!("failed to write checkpoint: {}", path.display()))?;
388
389 self.cleanup_old_snapshots().await?;
390
391 Ok(Some(metadata))
392 }
393
394 pub async fn list_snapshots(&self) -> Result<Vec<SnapshotMetadata>> {
395 if !self.enabled {
396 return Ok(Vec::new());
397 }
398 self.cleanup_old_snapshots().await?;
399 let snapshot_files = self.read_snapshot_files()?;
400 let mut snapshots = Vec::with_capacity(snapshot_files.len());
401 for (_, path) in snapshot_files {
402 let data = tokio::fs::read(&path)
403 .await
404 .with_context(|| format!("failed to read checkpoint: {}", path.display()))?;
405 let mut stored: StoredSnapshot = serde_json::from_slice(&data)
406 .with_context(|| format!("failed to parse checkpoint: {}", path.display()))?;
407 Self::hydrate_prompt_metadata(&mut stored);
408 snapshots.push(stored.metadata);
409 }
410 snapshots.sort_by(|a, b| b.turn_number.cmp(&a.turn_number));
411 Ok(snapshots)
412 }
413
414 pub async fn load_snapshot(&self, turn_number: usize) -> Result<Option<StoredSnapshot>> {
415 if !self.enabled {
416 return Ok(None);
417 }
418 let path = self.snapshot_path(turn_number);
419 if !tokio::fs::try_exists(&path).await.unwrap_or(false) {
420 return Ok(None);
421 }
422 let data = tokio::fs::read(&path)
423 .await
424 .with_context(|| format!("failed to read checkpoint: {}", path.display()))?;
425 let mut stored = serde_json::from_slice(&data)
426 .with_context(|| format!("failed to parse checkpoint: {}", path.display()))?;
427 Self::hydrate_prompt_metadata(&mut stored);
428 Ok(Some(stored))
429 }
430
431 pub async fn restore_snapshot(
432 &self,
433 turn_number: usize,
434 scope: RevertScope,
435 ) -> Result<Option<CheckpointRestore>> {
436 let Some(stored) = self.load_snapshot(turn_number).await? else {
437 return Ok(None);
438 };
439
440 if scope.includes_code() {
441 for snapshot in &stored.files {
442 let relative = Path::new(&snapshot.path);
443 let Some(sanitized) = sanitize_relative_path(relative) else {
444 continue;
445 };
446 let absolute = self.workspace.join(&sanitized);
447 if snapshot.deleted {
448 if tokio::fs::try_exists(&absolute).await.unwrap_or(false) {
449 tokio::fs::remove_file(&absolute).await.with_context(|| {
450 format!(
451 "failed to remove file during checkpoint restore: {}",
452 absolute.display()
453 )
454 })?;
455 }
456 continue;
457 }
458
459 if let Some(parent) = absolute.parent() {
460 ensure_dir_exists(parent).await.with_context(|| {
461 format!(
462 "failed to create directories for restore: {}",
463 parent.display()
464 )
465 })?;
466 }
467
468 let encoding = snapshot.encoding.unwrap_or(FileEncoding::Utf8);
469 let data = snapshot.data.as_deref().unwrap_or_default();
470 let bytes = Self::decode_file(encoding, data)?;
471 tokio::fs::write(&absolute, &bytes).await.with_context(|| {
472 format!("failed to write restored file: {}", absolute.display())
473 })?;
474 }
475 }
476
477 let conversation = if scope.includes_conversation() {
478 stored.conversation.clone()
479 } else {
480 Vec::new()
481 };
482
483 Ok(Some(CheckpointRestore {
484 metadata: stored.metadata,
485 conversation,
486 }))
487 }
488
489 pub async fn cleanup_old_snapshots(&self) -> Result<()> {
490 if !self.enabled {
491 return Ok(());
492 }
493
494 let mut entries = self.read_snapshot_files()?;
495
496 if let Some(cutoff) = self.retention_cutoff_secs()? {
497 let stale_entries = entries.clone();
498 for (_, path) in stale_entries {
499 let data = match tokio::fs::read(&path).await {
500 Ok(data) => data,
501 Err(err) => {
502 tracing::warn!(
503 path = %path.display(),
504 error = %err,
505 "Failed to read checkpoint"
506 );
507 continue;
508 }
509 };
510 let stored: StoredSnapshot = match serde_json::from_slice(&data) {
511 Ok(value) => value,
512 Err(err) => {
513 tracing::warn!(
514 path = %path.display(),
515 error = %err,
516 "Failed to parse checkpoint"
517 );
518 continue;
519 }
520 };
521 if stored.metadata.created_at <= cutoff
522 && let Err(err) = tokio::fs::remove_file(&path).await
523 {
524 tracing::warn!(
525 path = %path.display(),
526 error = %err,
527 "Failed to remove expired checkpoint"
528 );
529 }
530 }
531 entries = self.read_snapshot_files()?;
532 }
533
534 if self.max_snapshots == 0 || entries.len() <= self.max_snapshots {
535 return Ok(());
536 }
537
538 let excess = entries.len() - self.max_snapshots;
539 for (_, path) in entries.into_iter().take(excess) {
540 if let Err(err) = tokio::fs::remove_file(&path).await {
541 tracing::warn!(
542 path = %path.display(),
543 error = %err,
544 "Failed to remove old checkpoint"
545 );
546 }
547 }
548 Ok(())
549 }
550
551 fn retention_cutoff_secs(&self) -> Result<Option<u64>> {
552 let Some(days) = self.max_age_days else {
553 return Ok(None);
554 };
555
556 let now = Self::current_timestamp()?;
557 if days == 0 {
558 return Ok(Some(now));
559 }
560
561 let seconds = days.saturating_mul(SECONDS_PER_DAY);
562 let cutoff_instant = SystemTime::now()
563 .checked_sub(Duration::from_secs(seconds))
564 .unwrap_or(SystemTime::UNIX_EPOCH);
565 let cutoff = cutoff_instant
566 .duration_since(UNIX_EPOCH)
567 .context("system clock before UNIX_EPOCH")?
568 .as_secs();
569 Ok(Some(cutoff))
570 }
571
572 pub fn parse_revert_scope(value: &str) -> Option<RevertScope> {
573 match value.to_ascii_lowercase().as_str() {
574 "conversation" | "chat" => Some(RevertScope::Conversation),
575 "code" | "files" => Some(RevertScope::Code),
576 "both" | "full" => Some(RevertScope::Both),
577 _ => None,
578 }
579 }
580}
581
582#[derive(Debug, Clone)]
583pub struct CheckpointRestore {
584 pub metadata: SnapshotMetadata,
585 pub conversation: Vec<SessionMessage>,
586}
587
588#[cfg(test)]
589mod tests {
590 use super::*;
591 use tempfile::TempDir;
592
593 fn setup_manager() -> (TempDir, SnapshotManager) {
594 let dir = TempDir::new().expect("tempdir");
595 let workspace = dir.path().to_path_buf();
596 let manager =
597 SnapshotManager::new(SnapshotConfig::new(workspace.clone())).expect("manager");
598 (dir, manager)
599 }
600
601 #[tokio::test]
602 async fn create_and_list_snapshots() -> Result<()> {
603 let (_dir, manager) = setup_manager();
604 let mut conversation = Vec::new();
605 conversation.push(SessionMessage::new(
606 crate::llm::provider::MessageRole::User,
607 "Hello",
608 ));
609 let files = BTreeSet::new();
610 manager
611 .create_snapshot(1, "First turn", &conversation, &files, None, None)
612 .await?
613 .expect("metadata");
614 conversation.push(SessionMessage::new(
615 crate::llm::provider::MessageRole::Assistant,
616 "Hi",
617 ));
618 manager
619 .create_snapshot(2, "Second turn", &conversation, &files, None, None)
620 .await?
621 .expect("metadata");
622
623 let snapshots = manager.list_snapshots().await?;
624 assert_eq!(snapshots.len(), 2);
625 assert_eq!(snapshots[0].turn_number, 2);
626 assert_eq!(snapshots[1].turn_number, 1);
627 Ok(())
628 }
629
630 #[tokio::test]
631 async fn snapshot_restores_file_contents() -> Result<()> {
632 let (dir, manager) = setup_manager();
633 let workspace = dir.path();
634 let file_path = workspace.join("example.txt");
635 fs::write(&file_path, "v1")?;
636
637 let mut files = BTreeSet::new();
638 files.insert(PathBuf::from("example.txt"));
639 let conversation = vec![SessionMessage::new(
640 crate::llm::provider::MessageRole::User,
641 "edit example",
642 )];
643 manager
644 .create_snapshot(1, "save", &conversation, &files, None, None)
645 .await?
646 .expect("metadata");
647
648 fs::write(&file_path, "v2")?;
649 manager
650 .restore_snapshot(1, RevertScope::Code)
651 .await?
652 .expect("restore");
653 let restored = fs::read_to_string(&file_path)?;
654 assert_eq!(restored, "v1");
655 Ok(())
656 }
657
658 #[tokio::test]
659 async fn snapshot_handles_deleted_files() -> Result<()> {
660 let (dir, manager) = setup_manager();
661 let workspace = dir.path();
662 let file_path = workspace.join("remove.txt");
663 fs::write(&file_path, "data")?;
664
665 let mut files = BTreeSet::new();
666 files.insert(PathBuf::from("remove.txt"));
667 let conversation = vec![SessionMessage::new(
668 crate::llm::provider::MessageRole::User,
669 "remove",
670 )];
671 manager
672 .create_snapshot(1, "save", &conversation, &files, None, None)
673 .await?
674 .expect("metadata");
675
676 fs::remove_file(&file_path)?;
677 manager
678 .restore_snapshot(1, RevertScope::Code)
679 .await?
680 .expect("restore");
681 assert!(file_path.exists());
682 let content = fs::read_to_string(&file_path)?;
683 assert_eq!(content, "data");
684 Ok(())
685 }
686
687 #[tokio::test]
688 async fn cleanup_respects_limit() -> Result<()> {
689 let (_dir, manager) = setup_manager();
690 let conversation = vec![SessionMessage::new(
691 crate::llm::provider::MessageRole::User,
692 "hi",
693 )];
694 let files = BTreeSet::new();
695
696 for turn in 1..=5 {
697 manager
698 .create_snapshot(turn, "turn", &conversation, &files, None, None)
699 .await?
700 .expect("metadata");
701 }
702
703 let mut config = SnapshotConfig::new(manager.workspace.clone());
705 config.max_snapshots = 3;
706 let trimmed = SnapshotManager::new(config)?;
707 trimmed.cleanup_old_snapshots().await?;
708 let listed = trimmed.list_snapshots().await?;
709 assert_eq!(listed.len(), 3);
710 assert_eq!(listed[0].turn_number, 5);
711 assert_eq!(listed[2].turn_number, 3);
712 Ok(())
713 }
714
715 #[tokio::test]
716 async fn snapshot_normalizes_absolute_paths() -> Result<()> {
717 let (dir, manager) = setup_manager();
718 let workspace = dir.path();
719 let absolute = workspace.join("abs.txt");
720 fs::write(&absolute, "contents")?;
721
722 let mut files = BTreeSet::new();
723 files.insert(absolute.clone());
724 let conversation = vec![SessionMessage::new(
725 crate::llm::provider::MessageRole::User,
726 "absolute",
727 )];
728
729 manager
730 .create_snapshot(1, "abs", &conversation, &files, None, None)
731 .await?
732 .expect("metadata");
733
734 let stored = manager.load_snapshot(1).await?.expect("stored snapshot");
735 assert_eq!(stored.files.len(), 1);
736 assert_eq!(stored.files[0].path, "abs.txt");
737 assert!(!stored.files[0].deleted);
738 Ok(())
739 }
740
741 #[tokio::test]
742 async fn cleanup_removes_expired_snapshots() -> Result<()> {
743 let (_dir, manager) = setup_manager();
744 let conversation = vec![SessionMessage::new(
745 crate::llm::provider::MessageRole::User,
746 "cleanup",
747 )];
748 let files = BTreeSet::new();
749
750 manager
751 .create_snapshot(1, "old", &conversation, &files, None, None)
752 .await?
753 .expect("metadata");
754
755 let snapshot_path = manager.snapshot_path(1);
756 let mut stored: StoredSnapshot = serde_json::from_slice(&fs::read(&snapshot_path)?)?;
757 stored.metadata.created_at = 1;
758 let updated = serde_json::to_vec_pretty(&stored)?;
759 fs::write(&snapshot_path, updated)?;
760
761 let mut config = SnapshotConfig::new(manager.workspace.clone());
762 config.max_age_days = Some(1);
763 let janitor = SnapshotManager::new(config)?;
764 janitor.cleanup_old_snapshots().await?;
765
766 assert!(janitor.load_snapshot(1).await?.is_none());
767 Ok(())
768 }
769
770 #[tokio::test]
771 async fn snapshot_persists_prompt_metadata() -> Result<()> {
772 let (_dir, manager) = setup_manager();
773 let conversation = vec![
774 SessionMessage::new(
775 crate::llm::provider::MessageRole::User,
776 "Explain checkpointing",
777 ),
778 SessionMessage::new(
779 crate::llm::provider::MessageRole::Assistant,
780 "Working on it",
781 ),
782 ];
783
784 manager
785 .create_snapshot(
786 1,
787 "assistant reply",
788 &conversation,
789 &BTreeSet::new(),
790 Some("Explain checkpointing"),
791 Some(0),
792 )
793 .await?
794 .expect("metadata");
795
796 let stored = manager.load_snapshot(1).await?.expect("stored snapshot");
797 assert_eq!(
798 stored.metadata.prompt_text.as_deref(),
799 Some("Explain checkpointing")
800 );
801 assert_eq!(stored.metadata.prompt_message_index, Some(0));
802 assert_eq!(stored.metadata.description, "Explain checkpointing");
803 Ok(())
804 }
805
806 #[tokio::test]
807 async fn load_snapshot_hydrates_prompt_metadata_for_legacy_files() -> Result<()> {
808 let (_dir, manager) = setup_manager();
809 let stored = StoredSnapshot {
810 metadata: SnapshotMetadata {
811 id: "turn_1".to_string(),
812 turn_number: 1,
813 created_at: 1,
814 description: "legacy".to_string(),
815 message_count: 2,
816 file_count: 0,
817 prompt_text: None,
818 prompt_message_index: None,
819 },
820 conversation: vec![
821 SessionMessage::new(crate::llm::provider::MessageRole::User, "Legacy prompt"),
822 SessionMessage::new(crate::llm::provider::MessageRole::Assistant, "Legacy reply"),
823 ],
824 files: Vec::new(),
825 };
826 let path = manager.snapshot_path(1);
827 if let Some(parent) = path.parent() {
828 fs::create_dir_all(parent)?;
829 }
830 fs::write(path, serde_json::to_vec_pretty(&stored)?)?;
831
832 let loaded = manager.load_snapshot(1).await?.expect("loaded snapshot");
833 assert_eq!(
834 loaded.metadata.prompt_text.as_deref(),
835 Some("Legacy prompt")
836 );
837 assert_eq!(loaded.metadata.prompt_message_index, Some(0));
838 Ok(())
839 }
840
841 #[test]
842 fn parse_revert_scope_variants() {
843 assert_eq!(
844 SnapshotManager::parse_revert_scope("conversation"),
845 Some(RevertScope::Conversation)
846 );
847 assert_eq!(
848 SnapshotManager::parse_revert_scope("code"),
849 Some(RevertScope::Code)
850 );
851 assert_eq!(
852 SnapshotManager::parse_revert_scope("full"),
853 Some(RevertScope::Both)
854 );
855 assert_eq!(SnapshotManager::parse_revert_scope("unknown"), None);
856 }
857}