305 lines
9.5 KiB
Python
305 lines
9.5 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Rollback Manager
|
|
Allows undoing automation if it's not helpful
|
|
Creates backups and can restore to pre-automation state
|
|
"""
|
|
|
|
import json
|
|
import shutil
|
|
from pathlib import Path
|
|
from datetime import datetime
|
|
from typing import Dict, List, Optional
|
|
|
|
class RollbackManager:
|
|
"""Manages rollback of automation changes"""
|
|
|
|
def __init__(self, session_id: str):
|
|
self.session_id = session_id
|
|
self.backup_dir = Path(f".claude/meta-automation/backups/{session_id}")
|
|
self.manifest_path = self.backup_dir / "manifest.json"
|
|
self.backup_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
def create_backup(self, description: str = "Automation setup") -> str:
|
|
"""
|
|
Create backup before making changes
|
|
|
|
Args:
|
|
description: What this backup is for
|
|
|
|
Returns:
|
|
Backup ID
|
|
"""
|
|
backup_id = datetime.now().strftime("%Y%m%d_%H%M%S")
|
|
backup_path = self.backup_dir / backup_id
|
|
|
|
# Create backup manifest
|
|
manifest = {
|
|
'backup_id': backup_id,
|
|
'session_id': self.session_id,
|
|
'created_at': datetime.now().isoformat(),
|
|
'description': description,
|
|
'backed_up_files': [],
|
|
'created_files': [], # Files that didn't exist before
|
|
'can_rollback': True
|
|
}
|
|
|
|
# Save manifest
|
|
with open(self.manifest_path, 'w') as f:
|
|
json.dump(manifest, f, indent=2)
|
|
|
|
return backup_id
|
|
|
|
def track_file_creation(self, file_path: str):
|
|
"""
|
|
Track that a file was created by automation
|
|
|
|
Args:
|
|
file_path: Path to file that was created
|
|
"""
|
|
manifest = self._load_manifest()
|
|
if manifest:
|
|
if file_path not in manifest['created_files']:
|
|
manifest['created_files'].append(file_path)
|
|
self._save_manifest(manifest)
|
|
|
|
def backup_file_before_change(self, file_path: str):
|
|
"""
|
|
Backup a file before changing it
|
|
|
|
Args:
|
|
file_path: Path to file to backup
|
|
"""
|
|
manifest = self._load_manifest()
|
|
if not manifest:
|
|
return
|
|
|
|
source = Path(file_path)
|
|
if not source.exists():
|
|
return
|
|
|
|
# Create backup
|
|
backup_id = manifest['backup_id']
|
|
backup_path = self.backup_dir / backup_id
|
|
backup_path.mkdir(parents=True, exist_ok=True)
|
|
|
|
# Preserve directory structure in backup
|
|
rel_path = source.relative_to(Path.cwd()) if source.is_absolute() else source
|
|
dest = backup_path / rel_path
|
|
|
|
dest.parent.mkdir(parents=True, exist_ok=True)
|
|
shutil.copy2(source, dest)
|
|
|
|
# Track in manifest
|
|
if str(rel_path) not in manifest['backed_up_files']:
|
|
manifest['backed_up_files'].append(str(rel_path))
|
|
self._save_manifest(manifest)
|
|
|
|
def rollback(self) -> Dict:
|
|
"""
|
|
Rollback all automation changes
|
|
|
|
Returns:
|
|
Summary of what was rolled back
|
|
"""
|
|
manifest = self._load_manifest()
|
|
if not manifest:
|
|
return {
|
|
'success': False,
|
|
'message': 'No backup found for this session'
|
|
}
|
|
|
|
if not manifest['can_rollback']:
|
|
return {
|
|
'success': False,
|
|
'message': 'Rollback already performed or backup corrupted'
|
|
}
|
|
|
|
files_restored = []
|
|
files_deleted = []
|
|
errors = []
|
|
|
|
# Restore backed up files
|
|
backup_id = manifest['backup_id']
|
|
backup_path = self.backup_dir / backup_id
|
|
|
|
for file_path in manifest['backed_up_files']:
|
|
try:
|
|
source = backup_path / file_path
|
|
dest = Path(file_path)
|
|
|
|
if source.exists():
|
|
dest.parent.mkdir(parents=True, exist_ok=True)
|
|
shutil.copy2(source, dest)
|
|
files_restored.append(file_path)
|
|
else:
|
|
errors.append(f"Backup not found: {file_path}")
|
|
except Exception as e:
|
|
errors.append(f"Error restoring {file_path}: {str(e)}")
|
|
|
|
# Delete files that were created
|
|
for file_path in manifest['created_files']:
|
|
try:
|
|
path = Path(file_path)
|
|
if path.exists():
|
|
path.unlink()
|
|
files_deleted.append(file_path)
|
|
except Exception as e:
|
|
errors.append(f"Error deleting {file_path}: {str(e)}")
|
|
|
|
# Mark as rolled back
|
|
manifest['can_rollback'] = False
|
|
manifest['rolled_back_at'] = datetime.now().isoformat()
|
|
self._save_manifest(manifest)
|
|
|
|
return {
|
|
'success': len(errors) == 0,
|
|
'files_restored': files_restored,
|
|
'files_deleted': files_deleted,
|
|
'errors': errors,
|
|
'summary': f"Restored {len(files_restored)} files, deleted {len(files_deleted)} files"
|
|
}
|
|
|
|
def get_backup_info(self) -> Optional[Dict]:
|
|
"""Get information about current backup"""
|
|
manifest = self._load_manifest()
|
|
if not manifest:
|
|
return None
|
|
|
|
return {
|
|
'backup_id': manifest['backup_id'],
|
|
'created_at': manifest['created_at'],
|
|
'description': manifest['description'],
|
|
'backed_up_files_count': len(manifest['backed_up_files']),
|
|
'created_files_count': len(manifest['created_files']),
|
|
'can_rollback': manifest['can_rollback'],
|
|
'total_changes': len(manifest['backed_up_files']) + len(manifest['created_files'])
|
|
}
|
|
|
|
def preview_rollback(self) -> Dict:
|
|
"""
|
|
Preview what would be rolled back
|
|
|
|
Returns:
|
|
Details of what would happen
|
|
"""
|
|
manifest = self._load_manifest()
|
|
if not manifest:
|
|
return {
|
|
'can_rollback': False,
|
|
'message': 'No backup found'
|
|
}
|
|
|
|
return {
|
|
'can_rollback': manifest['can_rollback'],
|
|
'will_restore': manifest['backed_up_files'],
|
|
'will_delete': manifest['created_files'],
|
|
'total_changes': len(manifest['backed_up_files']) + len(manifest['created_files']),
|
|
'created_at': manifest['created_at'],
|
|
'description': manifest['description']
|
|
}
|
|
|
|
def _load_manifest(self) -> Optional[Dict]:
|
|
"""Load backup manifest"""
|
|
if not self.manifest_path.exists():
|
|
return None
|
|
|
|
try:
|
|
with open(self.manifest_path, 'r') as f:
|
|
return json.load(f)
|
|
except:
|
|
return None
|
|
|
|
def _save_manifest(self, manifest: Dict):
|
|
"""Save backup manifest"""
|
|
with open(self.manifest_path, 'w') as f:
|
|
json.dump(manifest, f, indent=2)
|
|
|
|
# Convenience wrapper for use in skills
|
|
class AutomationSnapshot:
|
|
"""
|
|
Context manager for creating automatic backups
|
|
|
|
Usage:
|
|
with AutomationSnapshot(session_id, "Adding security checks") as snapshot:
|
|
# Make changes
|
|
create_new_file("skill.md")
|
|
snapshot.track_creation("skill.md")
|
|
|
|
modify_file("existing.md")
|
|
snapshot.track_modification("existing.md")
|
|
|
|
# Automatic backup created, can rollback later
|
|
"""
|
|
|
|
def __init__(self, session_id: str, description: str):
|
|
self.manager = RollbackManager(session_id)
|
|
self.description = description
|
|
|
|
def __enter__(self):
|
|
self.manager.create_backup(self.description)
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
# Nothing to do on exit
|
|
pass
|
|
|
|
def track_creation(self, file_path: str):
|
|
"""Track file creation"""
|
|
self.manager.track_file_creation(file_path)
|
|
|
|
def track_modification(self, file_path: str):
|
|
"""Track file modification (backs up before change)"""
|
|
self.manager.backup_file_before_change(file_path)
|
|
|
|
# Example usage
|
|
if __name__ == '__main__':
|
|
import tempfile
|
|
import os
|
|
|
|
# Create test files
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
os.chdir(tmpdir)
|
|
|
|
# Create some test files
|
|
Path("existing.txt").write_text("original content")
|
|
|
|
manager = RollbackManager("test-session")
|
|
|
|
print("Creating backup...")
|
|
backup_id = manager.create_backup("Test automation setup")
|
|
|
|
# Simulate automation changes
|
|
print("\nMaking changes...")
|
|
manager.backup_file_before_change("existing.txt")
|
|
Path("existing.txt").write_text("modified content")
|
|
|
|
Path("new_skill.md").write_text("# New Skill")
|
|
manager.track_file_creation("new_skill.md")
|
|
|
|
Path("new_command.md").write_text("# New Command")
|
|
manager.track_file_creation("new_command.md")
|
|
|
|
# Show backup info
|
|
print("\nBackup info:")
|
|
info = manager.get_backup_info()
|
|
print(json.dumps(info, indent=2))
|
|
|
|
# Preview rollback
|
|
print("\nRollback preview:")
|
|
preview = manager.preview_rollback()
|
|
print(json.dumps(preview, indent=2))
|
|
|
|
# Perform rollback
|
|
print("\nPerforming rollback...")
|
|
result = manager.rollback()
|
|
print(json.dumps(result, indent=2))
|
|
|
|
# Check files
|
|
print("\nFiles after rollback:")
|
|
print(f"existing.txt exists: {Path('existing.txt').exists()}")
|
|
if Path("existing.txt").exists():
|
|
print(f"existing.txt content: {Path('existing.txt').read_text()}")
|
|
print(f"new_skill.md exists: {Path('new_skill.md').exists()}")
|
|
print(f"new_command.md exists: {Path('new_command.md').exists()}")
|