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

claude_mpm/services/infrastructure/memory_guardian.py

@@ -1,770 +0,0 @@
-"""Memory Guardian service for monitoring and managing Claude Code memory usage.
-
-This service monitors a subprocess (Claude Code) for memory consumption and
-performs automatic restarts when memory thresholds are exceeded.
-
-Design Principles:
-- Subprocess lifecycle management with graceful shutdown
-- Multi-threshold memory monitoring (warning, critical, emergency)
-- Platform-agnostic memory monitoring with fallbacks
-- Configurable restart policies with cooldown periods
-- State preservation hooks for future enhancement
-"""
-
-import asyncio
-import json
-import logging
-import os
-import platform
-import signal
-import subprocess
-import sys
-import time
-from dataclasses import dataclass, field
-from datetime import datetime, timedelta
-from enum import Enum
-from pathlib import Path
-from typing import Optional, Dict, Any, List, Callable, Tuple
-
-from claude_mpm.services.core.base import BaseService
-from claude_mpm.config.memory_guardian_config import (
-    MemoryGuardianConfig,
-    get_default_config
-)
-from claude_mpm.utils.platform_memory import (
-    get_process_memory,
-    get_system_memory,
-    check_memory_pressure,
-    MemoryInfo
-)
-
-
-class MemoryState(Enum):
-    """Memory usage state levels."""
-    NORMAL = "normal"
-    WARNING = "warning"
-    CRITICAL = "critical"
-    EMERGENCY = "emergency"
-
-
-class ProcessState(Enum):
-    """Process lifecycle states."""
-    NOT_STARTED = "not_started"
-    STARTING = "starting"
-    RUNNING = "running"
-    STOPPING = "stopping"
-    STOPPED = "stopped"
-    RESTARTING = "restarting"
-    FAILED = "failed"
-
-
-@dataclass
-class RestartAttempt:
-    """Record of a restart attempt."""
-    timestamp: float
-    reason: str
-    memory_mb: float
-    success: bool
-    
-    def to_dict(self) -> Dict[str, Any]:
-        """Convert to dictionary."""
-        return {
-            'timestamp': self.timestamp,
-            'timestamp_iso': datetime.fromtimestamp(self.timestamp).isoformat(),
-            'reason': self.reason,
-            'memory_mb': self.memory_mb,
-            'success': self.success
-        }
-
-
-@dataclass
-class MemoryStats:
-    """Memory usage statistics."""
-    current_mb: float = 0.0
-    peak_mb: float = 0.0
-    average_mb: float = 0.0
-    samples: int = 0
-    last_check: float = 0.0
-    state: MemoryState = MemoryState.NORMAL
-    
-    def update(self, memory_mb: float) -> None:
-        """Update statistics with new memory reading."""
-        self.current_mb = memory_mb
-        self.peak_mb = max(self.peak_mb, memory_mb)
-        
-        # Update running average
-        if self.samples == 0:
-            self.average_mb = memory_mb
-        else:
-            self.average_mb = ((self.average_mb * self.samples) + memory_mb) / (self.samples + 1)
-        
-        self.samples += 1
-        self.last_check = time.time()
-    
-    def to_dict(self) -> Dict[str, Any]:
-        """Convert to dictionary."""
-        return {
-            'current_mb': round(self.current_mb, 2),
-            'peak_mb': round(self.peak_mb, 2),
-            'average_mb': round(self.average_mb, 2),
-            'samples': self.samples,
-            'last_check': self.last_check,
-            'last_check_iso': datetime.fromtimestamp(self.last_check).isoformat() if self.last_check > 0 else None,
-            'state': self.state.value
-        }
-
-
-class MemoryGuardian(BaseService):
-    """Service for monitoring and managing subprocess memory usage."""
-    
-    def __init__(self, config: Optional[MemoryGuardianConfig] = None):
-        """Initialize Memory Guardian service.
-        
-        Args:
-            config: Configuration for memory monitoring and management
-        """
-        super().__init__("MemoryGuardian")
-        
-        # Configuration
-        self.config = config or get_default_config()
-        
-        # Validate configuration
-        issues = self.config.validate()
-        if issues:
-            for issue in issues:
-                self.log_warning(f"Configuration issue: {issue}")
-        
-        # Process management
-        self.process: Optional[subprocess.Popen] = None
-        self.process_state = ProcessState.NOT_STARTED
-        self.process_pid: Optional[int] = None
-        
-        # Memory monitoring
-        self.memory_stats = MemoryStats()
-        self.memory_state = MemoryState.NORMAL
-        
-        # Restart tracking
-        self.restart_attempts: List[RestartAttempt] = []
-        self.last_restart_time: float = 0.0
-        self.consecutive_failures: int = 0
-        
-        # Monitoring tasks
-        self.monitor_task: Optional[asyncio.Task] = None
-        self.monitoring = False
-        
-        # State preservation hooks (for future implementation)
-        self.state_save_hooks: List[Callable[[Dict[str, Any]], None]] = []
-        self.state_restore_hooks: List[Callable[[Dict[str, Any]], None]] = []
-        
-        # Statistics
-        self.start_time = time.time()
-        self.total_restarts = 0
-        self.total_uptime = 0.0
-        
-        self.log_info(f"Memory Guardian initialized with thresholds: "
-                     f"Warning={self.config.thresholds.warning}MB, "
-                     f"Critical={self.config.thresholds.critical}MB, "
-                     f"Emergency={self.config.thresholds.emergency}MB")
-    
-    async def initialize(self) -> bool:
-        """Initialize the Memory Guardian service.
-        
-        Returns:
-            True if initialization successful
-        """
-        try:
-            self.log_info("Initializing Memory Guardian service")
-            
-            # Load persisted state if available
-            if self.config.persist_state and self.config.state_file:
-                self._load_state()
-            
-            # Auto-start process if configured
-            if self.config.auto_start and self.config.enabled:
-                self.log_info("Auto-starting monitored process")
-                success = await self.start_process()
-                if not success:
-                    self.log_warning("Failed to auto-start process")
-            
-            # Start monitoring if enabled
-            if self.config.enabled:
-                self.start_monitoring()
-            
-            self._initialized = True
-            self.log_info("Memory Guardian service initialized successfully")
-            return True
-            
-        except Exception as e:
-            self.log_error(f"Failed to initialize Memory Guardian: {e}")
-            return False
-    
-    async def shutdown(self) -> None:
-        """Shutdown the Memory Guardian service gracefully."""
-        try:
-            self.log_info("Shutting down Memory Guardian service")
-            
-            # Stop monitoring
-            await self.stop_monitoring()
-            
-            # Save state if configured
-            if self.config.persist_state and self.config.state_file:
-                self._save_state()
-            
-            # Terminate process if running
-            if self.process and self.process_state == ProcessState.RUNNING:
-                await self.terminate_process()
-            
-            self._shutdown = True
-            self.log_info("Memory Guardian service shutdown complete")
-            
-        except Exception as e:
-            self.log_error(f"Error during Memory Guardian shutdown: {e}")
-    
-    async def start_process(self) -> bool:
-        """Start the monitored subprocess.
-        
-        Returns:
-            True if process started successfully
-        """
-        if self.process and self.process_state == ProcessState.RUNNING:
-            self.log_warning("Process is already running")
-            return True
-        
-        try:
-            self.log_info(f"Starting process: {' '.join(self.config.process_command)}")
-            self.process_state = ProcessState.STARTING
-            
-            # Prepare environment
-            env = os.environ.copy()
-            env.update(self.config.process_env)
-            
-            # Build command
-            cmd = self.config.process_command + self.config.process_args
-            
-            # Start subprocess
-            self.process = subprocess.Popen(
-                cmd,
-                env=env,
-                cwd=self.config.working_directory,
-                stdout=subprocess.PIPE,
-                stderr=subprocess.PIPE,
-                start_new_session=True  # Create new process group for clean termination
-            )
-            
-            self.process_pid = self.process.pid
-            self.process_state = ProcessState.RUNNING
-            
-            # Reset failure counter on successful start
-            self.consecutive_failures = 0
-            
-            self.log_info(f"Process started successfully with PID {self.process_pid}")
-            
-            # Give process time to initialize
-            await asyncio.sleep(2)
-            
-            # Check if process is still running
-            if self.process.poll() is not None:
-                self.log_error(f"Process exited immediately with code {self.process.returncode}")
-                self.process_state = ProcessState.FAILED
-                return False
-            
-            return True
-            
-        except FileNotFoundError:
-            self.log_error(f"Command not found: {self.config.process_command[0]}")
-            self.process_state = ProcessState.FAILED
-            return False
-        except Exception as e:
-            self.log_error(f"Failed to start process: {e}")
-            self.process_state = ProcessState.FAILED
-            return False
-    
-    async def restart_process(self, reason: str = "Manual restart") -> bool:
-        """Restart the monitored process with cooldown and retry logic.
-        
-        Args:
-            reason: Reason for restart
-            
-        Returns:
-            True if restart successful
-        """
-        self.log_info(f"Initiating process restart: {reason}")
-        
-        # Check restart attempts
-        if not self._can_restart():
-            self.log_error("Maximum restart attempts exceeded")
-            self.process_state = ProcessState.FAILED
-            return False
-        
-        # Apply cooldown if needed
-        cooldown = self._get_restart_cooldown()
-        if cooldown > 0:
-            self.log_info(f"Applying restart cooldown of {cooldown} seconds")
-            await asyncio.sleep(cooldown)
-        
-        # Record restart attempt
-        memory_mb = self.memory_stats.current_mb
-        self.process_state = ProcessState.RESTARTING
-        
-        # Save state before restart
-        await self._trigger_state_save()
-        
-        # Terminate existing process
-        if self.process:
-            await self.terminate_process()
-        
-        # Start new process
-        success = await self.start_process()
-        
-        # Record attempt
-        attempt = RestartAttempt(
-            timestamp=time.time(),
-            reason=reason,
-            memory_mb=memory_mb,
-            success=success
-        )
-        self.restart_attempts.append(attempt)
-        
-        if success:
-            self.total_restarts += 1
-            self.last_restart_time = time.time()
-            self.log_info("Process restarted successfully")
-            
-            # Restore state after restart
-            await self._trigger_state_restore()
-        else:
-            self.consecutive_failures += 1
-            self.log_error("Process restart failed")
-        
-        return success
-    
-    async def terminate_process(self, timeout: Optional[int] = None) -> bool:
-        """Terminate the monitored process gracefully with escalation.
-        
-        Args:
-            timeout: Override timeout for graceful shutdown
-            
-        Returns:
-            True if process terminated successfully
-        """
-        if not self.process:
-            return True
-        
-        timeout = timeout or self.config.restart_policy.graceful_timeout
-        
-        try:
-            self.log_info(f"Terminating process {self.process_pid}")
-            self.process_state = ProcessState.STOPPING
-            
-            # Try graceful termination first (SIGTERM)
-            if platform.system() != 'Windows':
-                self.process.terminate()
-            else:
-                # On Windows, terminate() is already forceful
-                self.process.terminate()
-            
-            # Wait for graceful shutdown
-            try:
-                self.log_debug(f"Waiting {timeout}s for graceful shutdown")
-                await asyncio.wait_for(
-                    asyncio.create_task(self._wait_for_process()),
-                    timeout=timeout
-                )
-                self.log_info("Process terminated gracefully")
-                
-            except asyncio.TimeoutError:
-                # Escalate to SIGKILL
-                self.log_warning("Graceful shutdown timeout, forcing termination")
-                
-                if platform.system() != 'Windows':
-                    self.process.kill()
-                else:
-                    # On Windows, use taskkill /F
-                    subprocess.run(
-                        ['taskkill', '/F', '/PID', str(self.process_pid)],
-                        capture_output=True
-                    )
-                
-                # Wait for forced termination
-                try:
-                    await asyncio.wait_for(
-                        asyncio.create_task(self._wait_for_process()),
-                        timeout=self.config.restart_policy.force_kill_timeout
-                    )
-                    self.log_info("Process terminated forcefully")
-                except asyncio.TimeoutError:
-                    self.log_error("Failed to terminate process")
-                    return False
-            
-            self.process = None
-            self.process_pid = None
-            self.process_state = ProcessState.STOPPED
-            return True
-            
-        except Exception as e:
-            self.log_error(f"Error terminating process: {e}")
-            return False
-    
-    async def _wait_for_process(self) -> None:
-        """Wait for process to exit."""
-        while self.process and self.process.poll() is None:
-            await asyncio.sleep(0.1)
-    
-    def get_memory_usage(self) -> Optional[float]:
-        """Get current memory usage of monitored process in MB.
-        
-        Returns:
-            Memory usage in MB or None if unable to determine
-        """
-        if not self.process or self.process_state != ProcessState.RUNNING:
-            return None
-        
-        try:
-            # Get memory info using platform utilities
-            mem_info = get_process_memory(self.process_pid)
-            if mem_info:
-                return mem_info.rss_mb
-            
-            self.log_warning(f"Unable to get memory info for PID {self.process_pid}")
-            return None
-            
-        except Exception as e:
-            self.log_error(f"Error getting memory usage: {e}")
-            return None
-    
-    async def monitor_memory(self) -> None:
-        """Check memory usage and take action if thresholds exceeded."""
-        if not self.process or self.process_state != ProcessState.RUNNING:
-            return
-        
-        # Check if process is still alive
-        if self.process.poll() is not None:
-            self.log_warning(f"Process exited with code {self.process.returncode}")
-            self.process_state = ProcessState.STOPPED
-            self.process = None
-            self.process_pid = None
-            
-            # Auto-restart if configured
-            if self.config.auto_start:
-                await self.restart_process("Process exited unexpectedly")
-            return
-        
-        # Get memory usage
-        memory_mb = self.get_memory_usage()
-        if memory_mb is None:
-            return
-        
-        # Update statistics
-        self.memory_stats.update(memory_mb)
-        
-        # Determine memory state
-        old_state = self.memory_state
-        
-        if memory_mb >= self.config.thresholds.emergency:
-            self.memory_state = MemoryState.EMERGENCY
-        elif memory_mb >= self.config.thresholds.critical:
-            self.memory_state = MemoryState.CRITICAL
-        elif memory_mb >= self.config.thresholds.warning:
-            self.memory_state = MemoryState.WARNING
-        else:
-            self.memory_state = MemoryState.NORMAL
-        
-        self.memory_stats.state = self.memory_state
-        
-        # Log state changes
-        if self.memory_state != old_state:
-            self.log_info(f"Memory state changed: {old_state.value} -> {self.memory_state.value} "
-                         f"(current: {memory_mb:.2f}MB)")
-        
-        # Take action based on state
-        if self.memory_state == MemoryState.EMERGENCY:
-            self.log_critical(f"Emergency memory threshold exceeded: {memory_mb:.2f}MB")
-            await self.restart_process(f"Emergency memory threshold exceeded ({memory_mb:.2f}MB)")
-            
-        elif self.memory_state == MemoryState.CRITICAL:
-            self.log_warning(f"Critical memory threshold exceeded: {memory_mb:.2f}MB")
-            # Check if we've been in critical state for too long
-            if self._should_restart_for_critical():
-                await self.restart_process(f"Sustained critical memory usage ({memory_mb:.2f}MB)")
-        
-        elif self.memory_state == MemoryState.WARNING:
-            self.log_debug(f"Warning memory threshold exceeded: {memory_mb:.2f}MB")
-        
-        # Log periodic summary
-        if self.config.monitoring.log_memory_stats:
-            if time.time() - self.memory_stats.last_check > self.config.monitoring.log_interval:
-                self._log_memory_summary()
-    
-    def start_monitoring(self) -> None:
-        """Start continuous memory monitoring."""
-        if self.monitoring:
-            self.log_warning("Monitoring is already active")
-            return
-        
-        self.monitoring = True
-        self.monitor_task = asyncio.create_task(self._monitoring_loop())
-        self.log_info("Started memory monitoring")
-    
-    async def stop_monitoring(self) -> None:
-        """Stop continuous memory monitoring."""
-        if not self.monitoring:
-            return
-        
-        self.monitoring = False
-        
-        if self.monitor_task:
-            self.monitor_task.cancel()
-            try:
-                await self.monitor_task
-            except asyncio.CancelledError:
-                pass
-            self.monitor_task = None
-        
-        self.log_info("Stopped memory monitoring")
-    
-    async def _monitoring_loop(self) -> None:
-        """Continuous monitoring loop."""
-        try:
-            while self.monitoring:
-                try:
-                    await self.monitor_memory()
-                    
-                    # Get check interval based on current state
-                    interval = self.config.monitoring.get_check_interval(
-                        self.memory_state.value
-                    )
-                    
-                    await asyncio.sleep(interval)
-                    
-                except Exception as e:
-                    self.log_error(f"Error in monitoring loop: {e}")
-                    await asyncio.sleep(5)  # Brief pause before retry
-                    
-        except asyncio.CancelledError:
-            self.log_debug("Monitoring loop cancelled")
-    
-    def _can_restart(self) -> bool:
-        """Check if restart is allowed based on policy.
-        
-        Returns:
-            True if restart is allowed
-        """
-        # Check max attempts
-        if self.config.restart_policy.max_attempts <= 0:
-            return True  # Unlimited restarts
-        
-        # Count recent attempts
-        window_start = time.time() - self.config.restart_policy.attempt_window
-        recent_attempts = [
-            a for a in self.restart_attempts
-            if a.timestamp >= window_start
-        ]
-        
-        return len(recent_attempts) < self.config.restart_policy.max_attempts
-    
-    def _get_restart_cooldown(self) -> int:
-        """Get cooldown period for next restart.
-        
-        Returns:
-            Cooldown period in seconds
-        """
-        if not self.restart_attempts:
-            return 0
-        
-        # Calculate based on consecutive failures
-        return self.config.restart_policy.get_cooldown(self.consecutive_failures + 1)
-    
-    def _should_restart_for_critical(self) -> bool:
-        """Determine if we should restart due to sustained critical memory.
-        
-        Returns:
-            True if restart should be triggered
-        """
-        # Check how long we've been in critical state
-        critical_duration = 60  # seconds
-        
-        # Look at recent memory samples
-        recent_samples = [
-            s for s in self.restart_attempts
-            if s.timestamp >= time.time() - critical_duration
-        ]
-        
-        # If we've been critical for the duration, restart
-        # This is a simplified check - could be enhanced
-        return self.memory_state == MemoryState.CRITICAL and len(recent_samples) == 0
-    
-    async def _trigger_state_save(self) -> None:
-        """Trigger state preservation hooks."""
-        if not self.state_save_hooks:
-            return
-        
-        state = self.get_state()
-        
-        for hook in self.state_save_hooks:
-            try:
-                hook(state)
-            except Exception as e:
-                self.log_error(f"State save hook failed: {e}")
-    
-    async def _trigger_state_restore(self) -> None:
-        """Trigger state restoration hooks."""
-        if not self.state_restore_hooks:
-            return
-        
-        state = self.get_state()
-        
-        for hook in self.state_restore_hooks:
-            try:
-                hook(state)
-            except Exception as e:
-                self.log_error(f"State restore hook failed: {e}")
-    
-    def _log_memory_summary(self) -> None:
-        """Log memory usage summary."""
-        uptime = time.time() - self.start_time
-        
-        self.log_info(
-            f"Memory Summary - "
-            f"Current: {self.memory_stats.current_mb:.2f}MB, "
-            f"Peak: {self.memory_stats.peak_mb:.2f}MB, "
-            f"Average: {self.memory_stats.average_mb:.2f}MB, "
-            f"State: {self.memory_state.value}, "
-            f"Restarts: {self.total_restarts}, "
-            f"Uptime: {uptime/3600:.2f}h"
-        )
-    
-    def _save_state(self) -> None:
-        """Save service state to file."""
-        if not self.config.state_file:
-            return
-        
-        try:
-            state = self.get_state()
-            state_path = Path(self.config.state_file)
-            state_path.parent.mkdir(parents=True, exist_ok=True)
-            
-            with open(state_path, 'w') as f:
-                json.dump(state, f, indent=2)
-            
-            self.log_debug(f"Saved state to {state_path}")
-            
-        except Exception as e:
-            self.log_error(f"Failed to save state: {e}")
-    
-    def _load_state(self) -> None:
-        """Load service state from file."""
-        if not self.config.state_file:
-            return
-        
-        try:
-            state_path = Path(self.config.state_file)
-            if not state_path.exists():
-                return
-            
-            with open(state_path, 'r') as f:
-                state = json.load(f)
-            
-            # Restore relevant state
-            self.total_restarts = state.get('total_restarts', 0)
-            self.memory_stats.peak_mb = state.get('memory_stats', {}).get('peak_mb', 0.0)
-            
-            # Restore restart attempts
-            attempts = state.get('restart_attempts', [])
-            for attempt_data in attempts:
-                attempt = RestartAttempt(
-                    timestamp=attempt_data['timestamp'],
-                    reason=attempt_data['reason'],
-                    memory_mb=attempt_data['memory_mb'],
-                    success=attempt_data['success']
-                )
-                self.restart_attempts.append(attempt)
-            
-            self.log_debug(f"Loaded state from {state_path}")
-            
-        except Exception as e:
-            self.log_error(f"Failed to load state: {e}")
-    
-    def add_state_save_hook(self, hook: Callable[[Dict[str, Any]], None]) -> None:
-        """Add a hook to be called before process restart.
-        
-        Args:
-            hook: Function to call with current state
-        """
-        self.state_save_hooks.append(hook)
-        self.log_debug(f"Added state save hook: {hook.__name__}")
-    
-    def add_state_restore_hook(self, hook: Callable[[Dict[str, Any]], None]) -> None:
-        """Add a hook to be called after process restart.
-        
-        Args:
-            hook: Function to call with saved state
-        """
-        self.state_restore_hooks.append(hook)
-        self.log_debug(f"Added state restore hook: {hook.__name__}")
-    
-    def get_state(self) -> Dict[str, Any]:
-        """Get current service state.
-        
-        Returns:
-            Dictionary containing service state
-        """
-        return {
-            'process_state': self.process_state.value,
-            'process_pid': self.process_pid,
-            'memory_state': self.memory_state.value,
-            'memory_stats': self.memory_stats.to_dict(),
-            'total_restarts': self.total_restarts,
-            'consecutive_failures': self.consecutive_failures,
-            'restart_attempts': [a.to_dict() for a in self.restart_attempts[-10:]],  # Last 10
-            'config': self.config.to_dict(),
-            'start_time': self.start_time,
-            'monitoring': self.monitoring
-        }
-    
-    def get_status(self) -> Dict[str, Any]:
-        """Get current service status.
-        
-        Returns:
-            Dictionary containing service status
-        """
-        uptime = time.time() - self.start_time if self.process else 0
-        
-        # Get system memory info
-        total_mem, available_mem = get_system_memory()
-        
-        return {
-            'enabled': self.config.enabled,
-            'process': {
-                'state': self.process_state.value,
-                'pid': self.process_pid,
-                'uptime_seconds': uptime,
-                'uptime_hours': uptime / 3600 if uptime > 0 else 0
-            },
-            'memory': {
-                'current_mb': self.memory_stats.current_mb,
-                'peak_mb': self.memory_stats.peak_mb,
-                'average_mb': self.memory_stats.average_mb,
-                'state': self.memory_state.value,
-                'thresholds': {
-                    'warning_mb': self.config.thresholds.warning,
-                    'critical_mb': self.config.thresholds.critical,
-                    'emergency_mb': self.config.thresholds.emergency
-                },
-                'system': {
-                    'total_mb': total_mem / (1024 * 1024) if total_mem > 0 else 0,
-                    'available_mb': available_mem / (1024 * 1024) if available_mem > 0 else 0,
-                    'pressure': check_memory_pressure()
-                }
-            },
-            'restarts': {
-                'total': self.total_restarts,
-                'consecutive_failures': self.consecutive_failures,
-                'can_restart': self._can_restart(),
-                'recent_attempts': [a.to_dict() for a in self.restart_attempts[-5:]]
-            },
-            'monitoring': {
-                'active': self.monitoring,
-                'check_interval': self.config.monitoring.get_check_interval(self.memory_state.value),
-                'samples': self.memory_stats.samples
-            }
-        }