Coverage for src / osiris_cli / checkpoint.py: 0%
181 statements
« prev ^ index » next coverage.py v7.13.0, created at 2025-12-27 17:41 +0200
1"""
2Session Checkpointing System
4Enables undo/rollback capability for AI changes.
5Users can rewind to previous states and experiment safely.
6"""
8import json
9import shutil
10from pathlib import Path
11from datetime import datetime
12from typing import List, Dict, Optional, Any
13from dataclasses import dataclass, asdict
15from .logger import get_logger
16from .config import CONFIG_DIR
18logger = get_logger()
20# Checkpoints directory
21CHECKPOINTS_DIR = CONFIG_DIR / "checkpoints"
22CHECKPOINTS_DIR.mkdir(parents=True, exist_ok=True)
25@dataclass
26class Checkpoint:
27 """A checkpoint of session state"""
28 id: str
29 session_id: str
30 timestamp: str
31 description: str
32 messages: List[Dict[str, Any]]
33 context_files: List[str]
34 file_snapshots: Dict[str, str] # file_path -> content
35 metadata: Dict[str, Any]
37 def to_dict(self) -> Dict:
38 return asdict(self)
40 @classmethod
41 def from_dict(cls, data: Dict) -> 'Checkpoint':
42 return cls(**data)
45class CheckpointManager:
46 """
47 Manage session checkpoints for undo/rollback.
49 Features:
50 - Auto-checkpoint after AI responses
51 - Manual checkpoint creation
52 - Restore to previous states
53 - File content snapshots
54 - Checkpoint diff visualization
55 """
57 def __init__(self, session_id: str, max_checkpoints: int = 20):
58 """
59 Initialize checkpoint manager.
61 Args:
62 session_id: Session identifier
63 max_checkpoints: Maximum number of checkpoints to keep
64 """
65 self.session_id = session_id
66 self.max_checkpoints = max_checkpoints
67 self.session_dir = CHECKPOINTS_DIR / session_id
68 self.session_dir.mkdir(parents=True, exist_ok=True)
70 self.checkpoints: List[Checkpoint] = []
71 self.load_checkpoints()
73 logger.info(f"Checkpoint manager initialized for session {session_id}")
75 def create_checkpoint(
76 self,
77 description: str,
78 messages: List[Dict],
79 context_files: List[str] = None,
80 files_to_snapshot: List[Path] = None,
81 metadata: Dict = None
82 ) -> Checkpoint:
83 """
84 Create a new checkpoint.
86 Args:
87 description: Checkpoint description
88 messages: Current conversation messages
89 context_files: Context files list
90 files_to_snapshot: Files to snapshot (save content)
91 metadata: Additional metadata
93 Returns:
94 Created checkpoint
95 """
96 checkpoint_id = f"cp_{datetime.now().strftime('%Y%m%d_%H%M%S_%f')}"
98 # Snapshot file contents
99 file_snapshots = {}
100 if files_to_snapshot:
101 for file_path in files_to_snapshot:
102 if file_path.exists():
103 try:
104 with open(file_path, 'r', encoding='utf-8') as f:
105 file_snapshots[str(file_path)] = f.read()
106 except Exception as e:
107 logger.warning(f"Failed to snapshot {file_path}: {e}")
109 checkpoint = Checkpoint(
110 id=checkpoint_id,
111 session_id=self.session_id,
112 timestamp=datetime.now().isoformat(),
113 description=description,
114 messages=messages.copy(),
115 context_files=context_files or [],
116 file_snapshots=file_snapshots,
117 metadata=metadata or {}
118 )
120 self.checkpoints.append(checkpoint)
121 self.save_checkpoint(checkpoint)
122 self.prune_checkpoints()
124 logger.info(f"Created checkpoint: {checkpoint_id} - {description}")
125 return checkpoint
127 def save_checkpoint(self, checkpoint: Checkpoint):
128 """Save checkpoint to disk"""
129 checkpoint_file = self.session_dir / f"{checkpoint.id}.json"
131 try:
132 with open(checkpoint_file, 'w') as f:
133 json.dump(checkpoint.to_dict(), f, indent=2)
134 logger.debug(f"Saved checkpoint: {checkpoint.id}")
135 except Exception as e:
136 logger.error(f"Failed to save checkpoint: {e}")
138 def load_checkpoints(self):
139 """Load all checkpoints for this session"""
140 self.checkpoints = []
142 for checkpoint_file in sorted(self.session_dir.glob("cp_*.json")):
143 try:
144 with open(checkpoint_file, 'r') as f:
145 data = json.load(f)
146 checkpoint = Checkpoint.from_dict(data)
147 self.checkpoints.append(checkpoint)
148 except Exception as e:
149 logger.error(f"Failed to load checkpoint {checkpoint_file}: {e}")
151 logger.debug(f"Loaded {len(self.checkpoints)} checkpoints")
153 def list_checkpoints(self) -> List[Dict[str, Any]]:
154 """Get list of checkpoints with summary info"""
155 return [
156 {
157 "id": cp.id,
158 "timestamp": cp.timestamp,
159 "description": cp.description,
160 "message_count": len(cp.messages),
161 "files_snapshotted": len(cp.file_snapshots)
162 }
163 for cp in sorted(self.checkpoints, key=lambda x: x.timestamp, reverse=True)
164 ]
166 def get_checkpoint(self, checkpoint_id: str) -> Optional[Checkpoint]:
167 """Get specific checkpoint by ID"""
168 for cp in self.checkpoints:
169 if cp.id == checkpoint_id:
170 return cp
171 return None
173 def restore_checkpoint(
174 self,
175 checkpoint_id: str,
176 restore_files: bool = True,
177 restore_messages: bool = True
178 ) -> Optional[Checkpoint]:
179 """
180 Restore session to checkpoint state.
182 Args:
183 checkpoint_id: Checkpoint ID to restore
184 restore_files: Whether to restore file contents
185 restore_messages: Whether to restore conversation
187 Returns:
188 Restored checkpoint or None
189 """
190 checkpoint = self.get_checkpoint(checkpoint_id)
192 if not checkpoint:
193 logger.error(f"Checkpoint not found: {checkpoint_id}")
194 return None
196 try:
197 # Restore file contents
198 if restore_files:
199 for file_path_str, content in checkpoint.file_snapshots.items():
200 file_path = Path(file_path_str)
202 # Create backup of current file
203 if file_path.exists():
204 backup_path = file_path.with_suffix(file_path.suffix + '.backup')
205 shutil.copy2(file_path, backup_path)
206 logger.debug(f"Backed up {file_path} to {backup_path}")
208 # Restore from checkpoint
209 file_path.parent.mkdir(parents=True, exist_ok=True)
210 with open(file_path, 'w', encoding='utf-8') as f:
211 f.write(content)
213 logger.info(f"Restored file: {file_path}")
215 logger.info(f"Restored checkpoint: {checkpoint_id}")
216 return checkpoint
218 except Exception as e:
219 logger.error(f"Failed to restore checkpoint: {e}", exc_info=True)
220 return None
222 def get_checkpoint_diff(
223 self,
224 checkpoint_id1: str,
225 checkpoint_id2: str
226 ) -> Dict[str, Any]:
227 """
228 Get diff between two checkpoints.
230 Args:
231 checkpoint_id1: First checkpoint ID
232 checkpoint_id2: Second checkpoint ID
234 Returns:
235 Diff summary
236 """
237 cp1 = self.get_checkpoint(checkpoint_id1)
238 cp2 = self.get_checkpoint(checkpoint_id2)
240 if not cp1 or not cp2:
241 return {}
243 # Message diff
244 message_diff = {
245 "added": len(cp2.messages) - len(cp1.messages),
246 "total_cp1": len(cp1.messages),
247 "total_cp2": len(cp2.messages)
248 }
250 # File diff
251 files_changed = []
252 for file_path in set(cp1.file_snapshots.keys()) | set(cp2.file_snapshots.keys()):
253 content1 = cp1.file_snapshots.get(file_path, "")
254 content2 = cp2.file_snapshots.get(file_path, "")
256 if content1 != content2:
257 files_changed.append({
258 "file": file_path,
259 "lines_cp1": len(content1.splitlines()),
260 "lines_cp2": len(content2.splitlines())
261 })
263 return {
264 "checkpoint1": {
265 "id": cp1.id,
266 "timestamp": cp1.timestamp,
267 "description": cp1.description
268 },
269 "checkpoint2": {
270 "id": cp2.id,
271 "timestamp": cp2.timestamp,
272 "description": cp2.description
273 },
274 "messages": message_diff,
275 "files_changed": files_changed
276 }
278 def prune_checkpoints(self):
279 """Remove old checkpoints beyond max limit"""
280 if len(self.checkpoints) <= self.max_checkpoints:
281 return
283 # Sort by timestamp
284 sorted_checkpoints = sorted(self.checkpoints, key=lambda x: x.timestamp)
286 # Remove oldest checkpoints
287 to_remove = sorted_checkpoints[:len(sorted_checkpoints) - self.max_checkpoints]
289 for checkpoint in to_remove:
290 checkpoint_file = self.session_dir / f"{checkpoint.id}.json"
292 try:
293 if checkpoint_file.exists():
294 checkpoint_file.unlink()
295 self.checkpoints.remove(checkpoint)
296 logger.debug(f"Pruned checkpoint: {checkpoint.id}")
297 except Exception as e:
298 logger.error(f"Failed to prune checkpoint {checkpoint.id}: {e}")
300 def delete_checkpoint(self, checkpoint_id: str) -> bool:
301 """Delete specific checkpoint"""
302 checkpoint = self.get_checkpoint(checkpoint_id)
304 if not checkpoint:
305 return False
307 try:
308 checkpoint_file = self.session_dir / f"{checkpoint.id}.json"
309 if checkpoint_file.exists():
310 checkpoint_file.unlink()
312 self.checkpoints.remove(checkpoint)
313 logger.info(f"Deleted checkpoint: {checkpoint_id}")
314 return True
316 except Exception as e:
317 logger.error(f"Failed to delete checkpoint: {e}")
318 return False
320 def export_checkpoint(
321 self,
322 checkpoint_id: str,
323 output_path: Path
324 ) -> bool:
325 """Export checkpoint to file"""
326 checkpoint = self.get_checkpoint(checkpoint_id)
328 if not checkpoint:
329 return False
331 try:
332 with open(output_path, 'w') as f:
333 json.dump(checkpoint.to_dict(), f, indent=2)
335 logger.info(f"Exported checkpoint to {output_path}")
336 return True
338 except Exception as e:
339 logger.error(f"Failed to export checkpoint: {e}")
340 return False
342 def import_checkpoint(self, input_path: Path) -> Optional[Checkpoint]:
343 """Import checkpoint from file"""
344 try:
345 with open(input_path, 'r') as f:
346 data = json.load(f)
347 checkpoint = Checkpoint.from_dict(data)
349 # Update session ID
350 checkpoint.session_id = self.session_id
352 self.checkpoints.append(checkpoint)
353 self.save_checkpoint(checkpoint)
355 logger.info(f"Imported checkpoint: {checkpoint.id}")
356 return checkpoint
358 except Exception as e:
359 logger.error(f"Failed to import checkpoint: {e}")
360 return None
362 def get_stats(self) -> Dict[str, Any]:
363 """Get checkpoint statistics"""
364 if not self.checkpoints:
365 return {
366 "total": 0,
367 "oldest": None,
368 "newest": None,
369 "total_size_mb": 0
370 }
372 # Calculate total size
373 total_size = 0
374 for cp_file in self.session_dir.glob("cp_*.json"):
375 total_size += cp_file.stat().st_size
377 sorted_cps = sorted(self.checkpoints, key=lambda x: x.timestamp)
379 return {
380 "total": len(self.checkpoints),
381 "oldest": sorted_cps[0].timestamp if sorted_cps else None,
382 "newest": sorted_cps[-1].timestamp if sorted_cps else None,
383 "total_size_mb": total_size / (1024 * 1024),
384 "max_checkpoints": self.max_checkpoints
385 }
388def create_auto_checkpoint(
389 checkpoint_manager: CheckpointManager,
390 description: str,
391 messages: List[Dict],
392 files_modified: List[Path] = None
393) -> Optional[Checkpoint]:
394 """
395 Create automatic checkpoint after AI response.
397 Args:
398 checkpoint_manager: CheckpointManager instance
399 description: Checkpoint description
400 messages: Current messages
401 files_modified: Files that were modified
403 Returns:
404 Created checkpoint or None
405 """
406 try:
407 checkpoint = checkpoint_manager.create_checkpoint(
408 description=f"Auto: {description}",
409 messages=messages,
410 files_to_snapshot=files_modified,
411 metadata={"auto": True}
412 )
413 return checkpoint
414 except Exception as e:
415 logger.error(f"Failed to create auto checkpoint: {e}")
416 return None