claude-mpm 3.9.9__py3-none-any.whl → 3.9.11__py3-none-any.whl

claude_mpm/services/infrastructure/state_manager.py

@@ -0,0 +1,774 @@
+"""State Manager service for capturing and restoring execution state.
+
+This service handles comprehensive state management for Claude Code,
+enabling seamless restarts with preserved context.
+
+Design Principles:
+- Atomic state operations with write-to-temp-then-rename
+- Automatic cleanup of old state files
+- Privacy-preserving (sanitizes sensitive data)
+- Platform-agnostic implementation
+- Graceful degradation on failures
+"""
+
+import asyncio
+import json
+import logging
+import os
+import pickle
+import platform
+import shutil
+import subprocess
+import sys
+import tempfile
+import time
+from datetime import datetime, timedelta
+from pathlib import Path
+from typing import Dict, Any, List, Optional, Tuple
+import uuid
+import gzip
+
+from claude_mpm.services.core.base import BaseService
+from claude_mpm.models.state_models import (
+    ProcessState,
+    ConversationState,
+    ConversationContext,
+    ProjectState,
+    RestartState,
+    CompleteState,
+    StateType
+)
+from claude_mpm.utils.platform_memory import get_process_memory
+
+
+class StateManager(BaseService):
+    """Service for managing Claude Code execution state across restarts."""
+    
+    def __init__(self, state_dir: Optional[Path] = None):
+        """Initialize State Manager service.
+        
+        Args:
+            state_dir: Directory for storing state files (default: ~/.claude-mpm/state)
+        """
+        super().__init__("StateManager")
+        
+        # State storage configuration
+        self.state_dir = state_dir or Path.home() / ".claude-mpm" / "state"
+        self.state_dir.mkdir(parents=True, exist_ok=True)
+        
+        # File naming
+        self.current_state_file = self.state_dir / "current_state.json"
+        self.compressed_state_file = self.state_dir / "current_state.json.gz"
+        
+        # Retention policy
+        self.retention_days = 7
+        self.max_state_files = 50
+        
+        # Claude conversation paths
+        self.claude_config_dir = Path.home() / ".claude"
+        self.claude_json_path = self.claude_config_dir / ".claude.json"
+        
+        # State tracking
+        self.current_state: Optional[CompleteState] = None
+        self.last_capture_time: float = 0
+        self.capture_cooldown = 5.0  # Minimum seconds between captures
+        
+        # Statistics
+        self.capture_count = 0
+        self.restore_count = 0
+        self.cleanup_count = 0
+        
+        self.log_info(f"State Manager initialized with state directory: {self.state_dir}")
+    
+    async def initialize(self) -> bool:
+        """Initialize the State Manager service.
+        
+        Returns:
+            True if initialization successful
+        """
+        try:
+            self.log_info("Initializing State Manager service")
+            
+            # Ensure state directory exists
+            self.state_dir.mkdir(parents=True, exist_ok=True)
+            
+            # Cleanup old states on startup
+            await self.cleanup_old_states()
+            
+            # Try to load existing state
+            loaded_state = await self.load_state()
+            if loaded_state:
+                self.log_info(f"Loaded existing state: {loaded_state.state_id}")
+                self.current_state = loaded_state
+            
+            self._initialized = True
+            self.log_info("State Manager service initialized successfully")
+            return True
+            
+        except Exception as e:
+            self.log_error(f"Failed to initialize State Manager: {e}")
+            return False
+    
+    async def shutdown(self) -> None:
+        """Shutdown the State Manager service gracefully."""
+        try:
+            self.log_info("Shutting down State Manager service")
+            
+            # Capture final state before shutdown
+            if self.current_state:
+                await self.persist_state(self.current_state)
+            
+            self._shutdown = True
+            self.log_info("State Manager service shutdown complete")
+            
+        except Exception as e:
+            self.log_error(f"Error during State Manager shutdown: {e}")
+    
+    async def capture_state(self, restart_reason: str = "Manual") -> Optional[CompleteState]:
+        """Capture current process and conversation state.
+        
+        Args:
+            restart_reason: Reason for the state capture/restart
+            
+        Returns:
+            CompleteState object or None if capture failed
+        """
+        try:
+            # Check cooldown
+            if time.time() - self.last_capture_time < self.capture_cooldown:
+                self.log_debug("State capture skipped due to cooldown")
+                return self.current_state
+            
+            self.log_info(f"Capturing state: {restart_reason}")
+            
+            # Capture process state
+            process_state = await self._capture_process_state()
+            
+            # Capture conversation state
+            conversation_state = await self._capture_conversation_state()
+            
+            # Capture project state
+            project_state = await self._capture_project_state()
+            
+            # Create restart state
+            restart_state = self._create_restart_state(restart_reason, process_state)
+            
+            # Create complete state
+            state = CompleteState(
+                process_state=process_state,
+                conversation_state=conversation_state,
+                project_state=project_state,
+                restart_state=restart_state
+            )
+            
+            # Validate state
+            issues = state.validate()
+            if issues:
+                for issue in issues:
+                    self.log_warning(f"State validation issue: {issue}")
+            
+            # Update tracking
+            self.current_state = state
+            self.last_capture_time = time.time()
+            self.capture_count += 1
+            
+            self.log_info(f"State captured successfully: {state.state_id}")
+            return state
+            
+        except Exception as e:
+            self.log_error(f"Failed to capture state: {e}")
+            return None
+    
+    async def restore_state(self, state: Optional[CompleteState] = None) -> bool:
+        """Restore state after process restart.
+        
+        Args:
+            state: State to restore (default: load from disk)
+            
+        Returns:
+            True if restoration successful
+        """
+        try:
+            # Load state if not provided
+            if state is None:
+                state = await self.load_state()
+                if state is None:
+                    self.log_warning("No state available to restore")
+                    return False
+            
+            self.log_info(f"Restoring state: {state.state_id}")
+            
+            # Restore working directory
+            if state.process_state.working_directory:
+                try:
+                    os.chdir(state.process_state.working_directory)
+                    self.log_debug(f"Restored working directory: {state.process_state.working_directory}")
+                except Exception as e:
+                    self.log_warning(f"Could not restore working directory: {e}")
+            
+            # Restore environment variables (non-sensitive only)
+            for key, value in state.process_state.environment.items():
+                if value != '***REDACTED***':
+                    os.environ[key] = value
+            
+            # Restore project-specific environment
+            if state.project_state.environment_vars:
+                for key, value in state.project_state.environment_vars.items():
+                    if value != '***REDACTED***':
+                        os.environ[key] = value
+            
+            # Log restoration summary
+            self.log_info(f"State restoration complete:")
+            self.log_info(f"  - Working directory: {state.process_state.working_directory}")
+            self.log_info(f"  - Git branch: {state.project_state.git_branch}")
+            self.log_info(f"  - Open files: {len(state.conversation_state.open_files)}")
+            self.log_info(f"  - Active conversation: {state.conversation_state.active_conversation_id}")
+            
+            self.restore_count += 1
+            return True
+            
+        except Exception as e:
+            self.log_error(f"Failed to restore state: {e}")
+            return False
+    
+    async def persist_state(self, state: CompleteState, compress: bool = True) -> bool:
+        """Save state to disk atomically.
+        
+        Args:
+            state: State to persist
+            compress: Whether to compress the state file
+            
+        Returns:
+            True if persistence successful
+        """
+        try:
+            self.log_debug(f"Persisting state: {state.state_id}")
+            
+            # Convert state to dictionary
+            state_dict = state.to_dict()
+            
+            # Create temporary file
+            temp_fd, temp_path = tempfile.mkstemp(
+                dir=self.state_dir,
+                suffix='.tmp',
+                prefix='state_'
+            )
+            
+            try:
+                # Write state to temporary file
+                if compress:
+                    # Write compressed JSON
+                    with gzip.open(temp_path, 'wt', encoding='utf-8') as f:
+                        json.dump(state_dict, f, indent=2)
+                    target_path = self.compressed_state_file
+                else:
+                    # Write plain JSON
+                    with os.fdopen(temp_fd, 'w') as f:
+                        json.dump(state_dict, f, indent=2)
+                    target_path = self.current_state_file
+                
+                # Atomic rename
+                temp_path_obj = Path(temp_path)
+                temp_path_obj.replace(target_path)
+                
+                # Also create timestamped backup
+                backup_name = f"state_{state.state_id}.json"
+                if compress:
+                    backup_name += ".gz"
+                backup_path = self.state_dir / backup_name
+                shutil.copy2(target_path, backup_path)
+                
+                self.log_debug(f"State persisted to {target_path}")
+                return True
+                
+            finally:
+                # Clean up temp file if it still exists
+                if Path(temp_path).exists():
+                    os.unlink(temp_path)
+                    
+        except Exception as e:
+            self.log_error(f"Failed to persist state: {e}")
+            return False
+    
+    async def load_state(self, state_file: Optional[Path] = None) -> Optional[CompleteState]:
+        """Load state from disk with validation.
+        
+        Args:
+            state_file: Path to state file (default: current_state.json[.gz])
+            
+        Returns:
+            CompleteState object or None if load failed
+        """
+        try:
+            # Determine which file to load
+            if state_file is None:
+                if self.compressed_state_file.exists():
+                    state_file = self.compressed_state_file
+                elif self.current_state_file.exists():
+                    state_file = self.current_state_file
+                else:
+                    self.log_debug("No state file found")
+                    return None
+            
+            self.log_debug(f"Loading state from {state_file}")
+            
+            # Read state file
+            if str(state_file).endswith('.gz'):
+                with gzip.open(state_file, 'rt', encoding='utf-8') as f:
+                    state_dict = json.load(f)
+            else:
+                with open(state_file, 'r') as f:
+                    state_dict = json.load(f)
+            
+            # Create state object
+            state = CompleteState.from_dict(state_dict)
+            
+            # Validate state
+            issues = state.validate()
+            if issues:
+                for issue in issues:
+                    self.log_warning(f"State validation issue: {issue}")
+            
+            self.log_info(f"Loaded state: {state.state_id}")
+            return state
+            
+        except Exception as e:
+            self.log_error(f"Failed to load state: {e}")
+            return None
+    
+    async def cleanup_old_states(self) -> int:
+        """Remove states older than retention period.
+        
+        Returns:
+            Number of files removed
+        """
+        try:
+            self.log_debug("Cleaning up old state files")
+            
+            removed_count = 0
+            cutoff_time = datetime.now() - timedelta(days=self.retention_days)
+            
+            # Find all state files
+            state_files = list(self.state_dir.glob("state_*.json*"))
+            
+            # Sort by modification time
+            state_files.sort(key=lambda f: f.stat().st_mtime)
+            
+            # Remove old files
+            for state_file in state_files:
+                try:
+                    # Check age
+                    file_time = datetime.fromtimestamp(state_file.stat().st_mtime)
+                    if file_time < cutoff_time:
+                        state_file.unlink()
+                        removed_count += 1
+                        self.log_debug(f"Removed old state file: {state_file.name}")
+                except Exception as e:
+                    self.log_warning(f"Could not remove state file {state_file}: {e}")
+            
+            # Also enforce max file limit
+            remaining_files = list(self.state_dir.glob("state_*.json*"))
+            if len(remaining_files) > self.max_state_files:
+                # Remove oldest files
+                remaining_files.sort(key=lambda f: f.stat().st_mtime)
+                for state_file in remaining_files[:-self.max_state_files]:
+                    try:
+                        state_file.unlink()
+                        removed_count += 1
+                        self.log_debug(f"Removed excess state file: {state_file.name}")
+                    except Exception as e:
+                        self.log_warning(f"Could not remove state file {state_file}: {e}")
+            
+            if removed_count > 0:
+                self.log_info(f"Cleaned up {removed_count} old state files")
+            
+            self.cleanup_count += removed_count
+            return removed_count
+            
+        except Exception as e:
+            self.log_error(f"Error during state cleanup: {e}")
+            return 0
+    
+    async def get_conversation_context(self) -> Optional[ConversationState]:
+        """Extract relevant Claude conversation data.
+        
+        Returns:
+            ConversationState object or None if extraction failed
+        """
+        try:
+            # Check if Claude config exists
+            if not self.claude_json_path.exists():
+                self.log_debug("Claude configuration not found")
+                return ConversationState(
+                    active_conversation_id=None,
+                    active_conversation=None,
+                    recent_conversations=[],
+                    total_conversations=0,
+                    total_storage_mb=0.0,
+                    preferences={},
+                    open_files=[],
+                    recent_files=[],
+                    pinned_files=[]
+                )
+            
+            # Get file size
+            file_size_mb = self.claude_json_path.stat().st_size / (1024 * 1024)
+            
+            # For large files, use streaming approach
+            if file_size_mb > 100:
+                self.log_warning(f"Large Claude config file ({file_size_mb:.2f}MB), using minimal extraction")
+                return self._extract_minimal_conversation_state()
+            
+            # Load Claude configuration
+            with open(self.claude_json_path, 'r') as f:
+                claude_data = json.load(f)
+            
+            # Extract conversation information
+            conversations = claude_data.get('conversations', [])
+            active_conv_id = claude_data.get('activeConversationId')
+            
+            # Find active conversation
+            active_conv = None
+            if active_conv_id:
+                for conv in conversations:
+                    if conv.get('id') == active_conv_id:
+                        active_conv = ConversationContext(
+                            conversation_id=conv.get('id', ''),
+                            title=conv.get('title', 'Untitled'),
+                            created_at=conv.get('createdAt', 0),
+                            updated_at=conv.get('updatedAt', 0),
+                            message_count=len(conv.get('messages', [])),
+                            total_tokens=conv.get('totalTokens', 0),
+                            max_tokens=conv.get('maxTokens', 100000),
+                            referenced_files=self._extract_referenced_files(conv),
+                            open_tabs=conv.get('openTabs', []),
+                            tags=conv.get('tags', []),
+                            is_active=True
+                        )
+                        break
+            
+            # Get recent conversations (last 5)
+            recent_convs = []
+            sorted_convs = sorted(
+                conversations,
+                key=lambda c: c.get('updatedAt', 0),
+                reverse=True
+            )[:5]
+            
+            for conv in sorted_convs:
+                if conv.get('id') != active_conv_id:
+                    recent_convs.append(ConversationContext(
+                        conversation_id=conv.get('id', ''),
+                        title=conv.get('title', 'Untitled'),
+                        created_at=conv.get('createdAt', 0),
+                        updated_at=conv.get('updatedAt', 0),
+                        message_count=len(conv.get('messages', [])),
+                        total_tokens=conv.get('totalTokens', 0),
+                        max_tokens=conv.get('maxTokens', 100000),
+                        referenced_files=self._extract_referenced_files(conv),
+                        open_tabs=conv.get('openTabs', []),
+                        tags=conv.get('tags', []),
+                        is_active=False
+                    ))
+            
+            # Extract preferences
+            preferences = claude_data.get('preferences', {})
+            
+            # Extract file state
+            open_files = claude_data.get('openFiles', [])
+            recent_files = claude_data.get('recentFiles', [])
+            pinned_files = claude_data.get('pinnedFiles', [])
+            
+            return ConversationState(
+                active_conversation_id=active_conv_id,
+                active_conversation=active_conv,
+                recent_conversations=recent_convs,
+                total_conversations=len(conversations),
+                total_storage_mb=file_size_mb,
+                preferences=preferences,
+                open_files=open_files,
+                recent_files=recent_files,
+                pinned_files=pinned_files
+            )
+            
+        except Exception as e:
+            self.log_error(f"Failed to extract conversation context: {e}")
+            return None
+    
+    async def _capture_process_state(self) -> ProcessState:
+        """Capture current process state."""
+        try:
+            pid = os.getpid()
+            
+            # Get memory info
+            mem_info = get_process_memory(pid)
+            memory_mb = mem_info.rss_mb if mem_info else 0.0
+            cpu_percent = 0.0  # Would need psutil for accurate CPU%
+            
+            # Get command line
+            if platform.system() != 'Windows':
+                try:
+                    with open(f'/proc/{pid}/cmdline', 'r') as f:
+                        cmdline = f.read().replace('\0', ' ').strip().split()
+                except:
+                    cmdline = sys.argv
+            else:
+                cmdline = sys.argv
+            
+            # Get open files (simplified - would need lsof or psutil for full list)
+            open_files = []
+            
+            return ProcessState(
+                pid=pid,
+                parent_pid=os.getppid(),
+                process_name=os.path.basename(sys.executable),
+                command=cmdline[:1] if cmdline else [],
+                args=cmdline[1:] if len(cmdline) > 1 else [],
+                working_directory=os.getcwd(),
+                environment=dict(os.environ),
+                memory_mb=memory_mb,
+                cpu_percent=cpu_percent,
+                open_files=open_files,
+                start_time=time.time() - time.process_time(),
+                capture_time=time.time()
+            )
+            
+        except Exception as e:
+            self.log_error(f"Error capturing process state: {e}")
+            # Return minimal state
+            return ProcessState(
+                pid=os.getpid(),
+                parent_pid=os.getppid(),
+                process_name="claude-mpm",
+                command=sys.argv[:1],
+                args=sys.argv[1:],
+                working_directory=os.getcwd(),
+                environment={},
+                memory_mb=0.0,
+                cpu_percent=0.0,
+                open_files=[],
+                start_time=time.time(),
+                capture_time=time.time()
+            )
+    
+    async def _capture_conversation_state(self) -> ConversationState:
+        """Capture conversation state from Claude."""
+        conversation_state = await self.get_conversation_context()
+        return conversation_state or ConversationState(
+            active_conversation_id=None,
+            active_conversation=None,
+            recent_conversations=[],
+            total_conversations=0,
+            total_storage_mb=0.0,
+            preferences={},
+            open_files=[],
+            recent_files=[],
+            pinned_files=[]
+        )
+    
+    async def _capture_project_state(self) -> ProjectState:
+        """Capture project and Git state."""
+        try:
+            project_path = os.getcwd()
+            project_name = os.path.basename(project_path)
+            
+            # Get Git information
+            git_branch = None
+            git_commit = None
+            git_status = {}
+            git_remotes = {}
+            
+            try:
+                # Get current branch
+                result = subprocess.run(
+                    ['git', 'rev-parse', '--abbrev-ref', 'HEAD'],
+                    capture_output=True,
+                    text=True,
+                    cwd=project_path
+                )
+                if result.returncode == 0:
+                    git_branch = result.stdout.strip()
+                
+                # Get current commit
+                result = subprocess.run(
+                    ['git', 'rev-parse', 'HEAD'],
+                    capture_output=True,
+                    text=True,
+                    cwd=project_path
+                )
+                if result.returncode == 0:
+                    git_commit = result.stdout.strip()[:8]
+                
+                # Get status
+                result = subprocess.run(
+                    ['git', 'status', '--porcelain'],
+                    capture_output=True,
+                    text=True,
+                    cwd=project_path
+                )
+                if result.returncode == 0:
+                    staged = []
+                    modified = []
+                    untracked = []
+                    
+                    for line in result.stdout.splitlines():
+                        if line.startswith('??'):
+                            untracked.append(line[3:])
+                        elif line.startswith(' M'):
+                            modified.append(line[3:])
+                        elif line.startswith('M '):
+                            staged.append(line[3:])
+                    
+                    git_status = {
+                        'staged': staged,
+                        'modified': modified,
+                        'untracked': untracked
+                    }
+                
+            except Exception as e:
+                self.log_debug(f"Could not get Git information: {e}")
+            
+            # Detect project type
+            project_type = self._detect_project_type(project_path)
+            
+            return ProjectState(
+                project_path=project_path,
+                project_name=project_name,
+                git_branch=git_branch,
+                git_commit=git_commit,
+                git_status=git_status,
+                git_remotes=git_remotes,
+                modified_files=git_status.get('modified', []),
+                open_editors=[],
+                breakpoints={},
+                project_type=project_type,
+                dependencies={},
+                environment_vars={},
+                last_build_status=None,
+                last_test_results=None
+            )
+            
+        except Exception as e:
+            self.log_error(f"Error capturing project state: {e}")
+            return ProjectState(
+                project_path=os.getcwd(),
+                project_name=os.path.basename(os.getcwd()),
+                git_branch=None,
+                git_commit=None,
+                git_status={},
+                git_remotes={},
+                modified_files=[],
+                open_editors=[],
+                breakpoints={},
+                project_type='unknown',
+                dependencies={},
+                environment_vars={},
+                last_build_status=None,
+                last_test_results=None
+            )
+    
+    def _create_restart_state(self, reason: str, process_state: ProcessState) -> RestartState:
+        """Create restart state information."""
+        return RestartState(
+            restart_id=str(uuid.uuid4()),
+            restart_count=self.capture_count,
+            timestamp=time.time(),
+            previous_uptime=time.time() - process_state.start_time,
+            reason=reason,
+            trigger='manual',  # Will be updated by caller
+            memory_mb=process_state.memory_mb,
+            memory_limit_mb=2048.0,  # Default limit
+            cpu_percent=process_state.cpu_percent,
+            error_type=None,
+            error_message=None,
+            error_traceback=None,
+            recovery_attempted=True,
+            recovery_successful=False,  # Will be updated after restore
+            data_preserved=['process', 'conversation', 'project']
+        )
+    
+    def _extract_referenced_files(self, conversation: Dict[str, Any]) -> List[str]:
+        """Extract file references from conversation."""
+        files = set()
+        
+        # Extract from messages
+        for message in conversation.get('messages', []):
+            content = message.get('content', '')
+            # Simple extraction - could be enhanced with regex
+            if isinstance(content, str):
+                # Look for file paths
+                import re
+                file_pattern = r'[\'"`]([^\'"`]+\.[a-zA-Z0-9]+)[\'"`]'
+                matches = re.findall(file_pattern, content)
+                files.update(matches)
+        
+        return list(files)[:100]  # Limit to prevent huge lists
+    
+    def _extract_minimal_conversation_state(self) -> ConversationState:
+        """Extract minimal conversation state for large files."""
+        try:
+            # Just get basic metadata without loading full file
+            file_size_mb = self.claude_json_path.stat().st_size / (1024 * 1024)
+            
+            return ConversationState(
+                active_conversation_id="large_file",
+                active_conversation=None,
+                recent_conversations=[],
+                total_conversations=-1,  # Unknown
+                total_storage_mb=file_size_mb,
+                preferences={},
+                open_files=[],
+                recent_files=[],
+                pinned_files=[]
+            )
+        except:
+            return ConversationState(
+                active_conversation_id=None,
+                active_conversation=None,
+                recent_conversations=[],
+                total_conversations=0,
+                total_storage_mb=0.0,
+                preferences={},
+                open_files=[],
+                recent_files=[],
+                pinned_files=[]
+            )
+    
+    def _detect_project_type(self, project_path: str) -> str:
+        """Detect project type from files present."""
+        path = Path(project_path)
+        
+        if (path / 'pyproject.toml').exists() or (path / 'setup.py').exists():
+            return 'python'
+        elif (path / 'package.json').exists():
+            return 'node'
+        elif (path / 'go.mod').exists():
+            return 'go'
+        elif (path / 'Cargo.toml').exists():
+            return 'rust'
+        elif (path / 'pom.xml').exists():
+            return 'java'
+        else:
+            return 'unknown'
+    
+    def get_statistics(self) -> Dict[str, Any]:
+        """Get state manager statistics.
+        
+        Returns:
+            Dictionary containing statistics
+        """
+        state_files = list(self.state_dir.glob("state_*.json*"))
+        total_size_mb = sum(f.stat().st_size for f in state_files) / (1024 * 1024)
+        
+        return {
+            'capture_count': self.capture_count,
+            'restore_count': self.restore_count,
+            'cleanup_count': self.cleanup_count,
+            'state_files': len(state_files),
+            'total_size_mb': round(total_size_mb, 2),
+            'state_directory': str(self.state_dir),
+            'current_state_id': self.current_state.state_id if self.current_state else None,
+            'last_capture_time': self.last_capture_time,
+            'retention_days': self.retention_days,
+            'max_state_files': self.max_state_files
+        }