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

claude_mpm/services/infrastructure/memory_guardian.py

@@ -37,6 +37,10 @@ from claude_mpm.utils.platform_memory import (
     check_memory_pressure,
     MemoryInfo
 )
+from claude_mpm.services.infrastructure.state_manager import StateManager
+from claude_mpm.services.infrastructure.restart_protection import RestartProtection
+from claude_mpm.services.infrastructure.health_monitor import HealthMonitor
+from claude_mpm.services.infrastructure.graceful_degradation import GracefulDegradation
 
 
 class MemoryState(Enum):
@@ -156,6 +160,14 @@ class MemoryGuardian(BaseService):
         self.state_save_hooks: List[Callable[[Dict[str, Any]], None]] = []
         self.state_restore_hooks: List[Callable[[Dict[str, Any]], None]] = []
         
+        # State manager integration
+        self.state_manager: Optional[StateManager] = None
+        
+        # Safety services integration
+        self.restart_protection: Optional[RestartProtection] = None
+        self.health_monitor: Optional[HealthMonitor] = None
+        self.graceful_degradation: Optional[GracefulDegradation] = None
+        
         # Statistics
         self.start_time = time.time()
         self.total_restarts = 0
@@ -175,6 +187,13 @@ class MemoryGuardian(BaseService):
         try:
             self.log_info("Initializing Memory Guardian service")
             
+            # Initialize state manager
+            self.state_manager = StateManager()
+            await self.state_manager.initialize()
+            
+            # Initialize safety services
+            await self._initialize_safety_services()
+            
             # Load persisted state if available
             if self.config.persist_state and self.config.state_file:
                 self._load_state()
@@ -210,6 +229,18 @@ class MemoryGuardian(BaseService):
             if self.config.persist_state and self.config.state_file:
                 self._save_state()
             
+            # Shutdown state manager
+            if self.state_manager:
+                await self.state_manager.shutdown()
+            
+            # Shutdown safety services
+            if self.restart_protection:
+                await self.restart_protection.shutdown()
+            if self.health_monitor:
+                await self.health_monitor.shutdown()
+            if self.graceful_degradation:
+                await self.graceful_degradation.shutdown()
+            
             # Terminate process if running
             if self.process and self.process_state == ProcessState.RUNNING:
                 await self.terminate_process()
@@ -254,6 +285,10 @@ class MemoryGuardian(BaseService):
             self.process_pid = self.process.pid
             self.process_state = ProcessState.RUNNING
             
+            # Update health monitor with new process
+            if self.health_monitor:
+                self.health_monitor.set_monitored_process(self.process_pid)
+            
             # Reset failure counter on successful start
             self.consecutive_failures = 0
             
@@ -290,24 +325,56 @@ class MemoryGuardian(BaseService):
         """
         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)
+        # Use restart protection if available
+        if self.restart_protection:
+            allowed, protection_reason = self.restart_protection.should_allow_restart(
+                self.memory_stats.current_mb
+            )
+            if not allowed:
+                self.log_error(f"Restart blocked by protection: {protection_reason}")
+                
+                # Trigger graceful degradation if available
+                if self.graceful_degradation:
+                    await self.graceful_degradation.disable_feature(
+                        "automated_monitoring",
+                        f"Restart protection triggered: {protection_reason}"
+                    )
+                
+                self.process_state = ProcessState.FAILED
+                return False
+            
+            # Get backoff from restart protection
+            backoff = self.restart_protection.get_backoff_seconds(
+                self.restart_protection.statistics.consecutive_failures + 1
+            )
+            if backoff > 0:
+                self.log_info(f"Applying restart backoff of {backoff:.1f} seconds")
+                await asyncio.sleep(backoff)
+        else:
+            # Fallback to original logic
+            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()
+        # Save state before restart using StateManager
+        if self.state_manager:
+            state = await self.state_manager.capture_state(restart_reason=reason)
+            if state:
+                await self.state_manager.persist_state(state)
+        else:
+            # Fallback to hook-based preservation
+            await self._trigger_state_save()
         
         # Terminate existing process
         if self.process:
@@ -325,13 +392,29 @@ class MemoryGuardian(BaseService):
         )
         self.restart_attempts.append(attempt)
         
+        # Record in restart protection service
+        if self.restart_protection:
+            backoff_applied = self.restart_protection.get_backoff_seconds(
+                self.restart_protection.statistics.consecutive_failures + 1
+            ) if not success else 0
+            self.restart_protection.record_restart(
+                reason=reason,
+                memory_mb=memory_mb,
+                success=success,
+                backoff_applied=backoff_applied
+            )
+        
         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()
+            # Restore state after restart using StateManager
+            if self.state_manager:
+                await self.state_manager.restore_state()
+            else:
+                # Fallback to hook-based restoration
+                await self._trigger_state_restore()
         else:
             self.consecutive_failures += 1
             self.log_error("Process restart failed")
@@ -457,6 +540,10 @@ class MemoryGuardian(BaseService):
         # Update statistics
         self.memory_stats.update(memory_mb)
         
+        # Record memory sample for trend analysis
+        if self.restart_protection:
+            self.restart_protection.record_memory_sample(memory_mb)
+        
         # Determine memory state
         old_state = self.memory_state
         
@@ -702,6 +789,15 @@ class MemoryGuardian(BaseService):
         self.state_restore_hooks.append(hook)
         self.log_debug(f"Added state restore hook: {hook.__name__}")
     
+    def set_state_manager(self, state_manager: StateManager) -> None:
+        """Set the state manager for state preservation.
+        
+        Args:
+            state_manager: StateManager instance to use
+        """
+        self.state_manager = state_manager
+        self.log_info("State manager configured for Memory Guardian")
+    
     def get_state(self) -> Dict[str, Any]:
         """Get current service state.
         
@@ -721,6 +817,84 @@ class MemoryGuardian(BaseService):
             'monitoring': self.monitoring
         }
     
+    async def _initialize_safety_services(self) -> None:
+        """Initialize safety and protection services."""
+        try:
+            # Initialize restart protection
+            self.restart_protection = RestartProtection(
+                max_restarts_per_hour=5,
+                max_consecutive_failures=3,
+                base_backoff_seconds=1,
+                max_backoff_seconds=60,
+                state_file=Path(self.config.state_file).parent / "restart_protection.json" if self.config.state_file else None
+            )
+            await self.restart_protection.initialize()
+            self.log_info("Restart protection service initialized")
+            
+            # Initialize health monitor
+            self.health_monitor = HealthMonitor(
+                cpu_threshold_percent=80,
+                memory_threshold_percent=90,
+                disk_threshold_percent=90,
+                min_disk_space_gb=1.0,
+                check_interval_seconds=30
+            )
+            await self.health_monitor.initialize()
+            self.log_info("Health monitor service initialized")
+            
+            # Initialize graceful degradation
+            self.graceful_degradation = GracefulDegradation(
+                enable_notifications=True,
+                log_degradation_events=True,
+                state_file=Path(self.config.state_file).parent / "degradation.json" if self.config.state_file else None
+            )
+            await self.graceful_degradation.initialize()
+            self.log_info("Graceful degradation service initialized")
+            
+            # Check initial health
+            valid, message = await self.health_monitor.validate_before_start()
+            if not valid:
+                self.log_warning(f"System health check warning: {message}")
+                await self.graceful_degradation.degrade_feature(
+                    "automated_monitoring",
+                    message,
+                    "reduced monitoring frequency"
+                )
+            
+        except Exception as e:
+            self.log_error(f"Failed to initialize safety services: {e}")
+            # Continue without safety services - graceful degradation
+            self.log_warning("Continuing without safety services")
+    
+    def set_restart_protection(self, restart_protection: RestartProtection) -> None:
+        """Set the restart protection service.
+        
+        Args:
+            restart_protection: RestartProtection instance to use
+        """
+        self.restart_protection = restart_protection
+        self.log_info("Restart protection configured for Memory Guardian")
+    
+    def set_health_monitor(self, health_monitor: HealthMonitor) -> None:
+        """Set the health monitor service.
+        
+        Args:
+            health_monitor: HealthMonitor instance to use
+        """
+        self.health_monitor = health_monitor
+        if self.process_pid:
+            self.health_monitor.set_monitored_process(self.process_pid)
+        self.log_info("Health monitor configured for Memory Guardian")
+    
+    def set_graceful_degradation(self, graceful_degradation: GracefulDegradation) -> None:
+        """Set the graceful degradation service.
+        
+        Args:
+            graceful_degradation: GracefulDegradation instance to use
+        """
+        self.graceful_degradation = graceful_degradation
+        self.log_info("Graceful degradation configured for Memory Guardian")
+    
     def get_status(self) -> Dict[str, Any]:
         """Get current service status.