agent_os_kernel 3.1.0__py3-none-any.whl

agent_control_plane/agent_hibernation.py

@@ -0,0 +1,394 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""
+Agent Hibernation - State-to-Disk Serialization
+
+Feature: "Agent Hibernation (state-to-disk)"
+Problem: Agents sitting idle in memory cost money (if hosted) or RAM.
+Solution: Serialize the entire agent state (including caas pointer) to a JSON/Pickle file
+         and kill the process. Wake it up only when amb receives a message for it.
+Result: Scale by Subtraction - Removes the need for "always-on" servers. Serverless Agents.
+
+This module provides the infrastructure to hibernate idle agents by serializing
+their complete state to disk and restoring them when needed.
+"""
+
+from typing import Any, Dict, Optional, List
+from dataclasses import dataclass, field, asdict
+from datetime import datetime, timedelta
+from enum import Enum
+import json
+import hmac
+import hashlib
+import os
+import logging
+from pathlib import Path
+
+
+class HibernationFormat(Enum):
+    """Format for serializing agent state"""
+    JSON = "json"
+    PICKLE = "pickle"
+
+
+class AgentState(Enum):
+    """State of an agent in the hibernation lifecycle"""
+    ACTIVE = "active"
+    HIBERNATING = "hibernating"
+    HIBERNATED = "hibernated"
+    WAKING = "waking"
+
+
+@dataclass
+class HibernatedAgentMetadata:
+    """Metadata about a hibernated agent"""
+    agent_id: str
+    session_id: str
+    hibernated_at: datetime
+    state_file_path: str
+    format: HibernationFormat
+    state_size_bytes: int
+    last_activity: datetime
+    context_pointer: Optional[str] = None  # caas pointer
+    metadata: Dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class HibernationConfig:
+    """Configuration for agent hibernation"""
+    enabled: bool = True
+    idle_timeout_seconds: int = 300  # 5 minutes default
+    storage_path: str = "/tmp/agent_hibernation"
+    format: HibernationFormat = HibernationFormat.JSON
+    max_hibernated_agents: int = 1000
+    auto_cleanup_days: int = 7  # Clean up hibernated agents after 7 days
+    compress: bool = False
+
+
+class HibernationManager:
+    """
+    Manages agent hibernation and wake-up operations.
+    
+    This class handles:
+    - Serialization of complete agent state to disk
+    - Deserialization and restoration of agent state
+    - Tracking of hibernated agents
+    - Automatic cleanup of old hibernated states
+    """
+    
+    def __init__(self, config: Optional[HibernationConfig] = None):
+        """
+        Initialize the hibernation manager.
+        
+        Args:
+            config: Configuration for hibernation behavior
+        """
+        self.config = config or HibernationConfig()
+        self.logger = logging.getLogger("HibernationManager")
+        
+        # Create storage directory if it doesn't exist
+        Path(self.config.storage_path).mkdir(parents=True, exist_ok=True)
+        
+        # HMAC key for pickle integrity verification — generated per instance.
+        # In production, this should be persisted securely (e.g., Key Vault).
+        self._hmac_key = os.urandom(32)
+        
+        # Create storage directory if it doesn't exist
+        Path(self.config.storage_path).mkdir(parents=True, exist_ok=True)
+        
+        # Track hibernated agents
+        self.hibernated_agents: Dict[str, HibernatedAgentMetadata] = {}
+        
+        # Track agent activity for idle detection
+        self.agent_activity: Dict[str, datetime] = {}
+        
+        # Track agent states
+        self.agent_states: Dict[str, AgentState] = {}
+        
+        self.logger.info(f"HibernationManager initialized with storage at {self.config.storage_path}")
+    
+    def serialize_agent_state(
+        self,
+        agent_id: str,
+        agent_context: Any,
+        caas_pointer: Optional[str] = None,
+        additional_state: Optional[Dict[str, Any]] = None
+    ) -> Dict[str, Any]:
+        """
+        Serialize agent state to a dictionary.
+        
+        Args:
+            agent_id: Unique identifier for the agent
+            agent_context: AgentContext object containing agent session data
+            caas_pointer: Optional pointer to context in caas (Context-as-a-Service)
+            additional_state: Optional additional state to serialize
+            
+        Returns:
+            Dictionary containing serialized agent state
+        """
+        state = {
+            "agent_id": agent_id,
+            "session_id": agent_context.session_id,
+            "created_at": agent_context.created_at.isoformat(),
+            "permissions": {
+                str(k): v.value for k, v in agent_context.permissions.items()
+            },
+            "metadata": agent_context.metadata,
+            "caas_pointer": caas_pointer,
+            "hibernated_at": datetime.now().isoformat(),
+            "additional_state": additional_state or {}
+        }
+        
+        return state
+    
+    def deserialize_agent_state(self, state: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        Deserialize agent state from a dictionary.
+        
+        Args:
+            state: Dictionary containing serialized agent state
+            
+        Returns:
+            Dictionary with deserialized state components
+        """
+        # Convert ISO format strings back to datetime
+        if "created_at" in state:
+            state["created_at"] = datetime.fromisoformat(state["created_at"])
+        if "hibernated_at" in state:
+            state["hibernated_at"] = datetime.fromisoformat(state["hibernated_at"])
+        
+        return state
+    
+    def hibernate_agent(
+        self,
+        agent_id: str,
+        agent_context: Any,
+        caas_pointer: Optional[str] = None,
+        additional_state: Optional[Dict[str, Any]] = None
+    ) -> HibernatedAgentMetadata:
+        """
+        Hibernate an agent by saving its state to disk.
+        
+        Args:
+            agent_id: Unique identifier for the agent
+            agent_context: AgentContext object to serialize
+            caas_pointer: Optional pointer to context in caas
+            additional_state: Optional additional state to save
+            
+        Returns:
+            Metadata about the hibernated agent
+        """
+        self.logger.info(f"Hibernating agent {agent_id}")
+        
+        # Update agent state
+        self.agent_states[agent_id] = AgentState.HIBERNATING
+        
+        # Serialize agent state
+        state = self.serialize_agent_state(agent_id, agent_context, caas_pointer, additional_state)
+        
+        # Determine file path
+        file_name = f"{agent_id}_{agent_context.session_id}.{self.config.format.value}"
+        file_path = os.path.join(self.config.storage_path, file_name)
+        
+        # Save to disk
+        try:
+            if self.config.format == HibernationFormat.JSON:
+                with open(file_path, 'w') as f:
+                    json.dump(state, f, indent=2)
+            else:  # PICKLE format — now uses JSON internally + HMAC signature
+                raw = json.dumps(state).encode('utf-8')
+                sig = hmac.new(self._hmac_key, raw, hashlib.sha256).hexdigest()
+                with open(file_path, 'wb') as f:
+                    f.write(raw)
+                with open(file_path + ".sig", 'w', encoding='utf-8') as f:
+                    f.write(sig)
+            
+            # Get file size
+            state_size = os.path.getsize(file_path)
+            
+            # Create metadata
+            metadata = HibernatedAgentMetadata(
+                agent_id=agent_id,
+                session_id=agent_context.session_id,
+                hibernated_at=datetime.now(),
+                state_file_path=file_path,
+                format=self.config.format,
+                state_size_bytes=state_size,
+                last_activity=self.agent_activity.get(agent_id, datetime.now()),
+                context_pointer=caas_pointer
+            )
+            
+            # Track hibernated agent
+            self.hibernated_agents[agent_id] = metadata
+            self.agent_states[agent_id] = AgentState.HIBERNATED
+            
+            self.logger.info(f"Agent {agent_id} hibernated successfully. State saved to {file_path} ({state_size} bytes)")
+            
+            return metadata
+            
+        except Exception as e:
+            self.logger.error(f"Failed to hibernate agent {agent_id}: {e}")
+            self.agent_states[agent_id] = AgentState.ACTIVE
+            raise
+    
+    def wake_agent(self, agent_id: str) -> Dict[str, Any]:
+        """
+        Wake up a hibernated agent by restoring its state from disk.
+        
+        Args:
+            agent_id: Unique identifier for the agent to wake
+            
+        Returns:
+            Dictionary containing the restored agent state
+        """
+        if agent_id not in self.hibernated_agents:
+            raise ValueError(f"Agent {agent_id} is not hibernated")
+        
+        self.logger.info(f"Waking agent {agent_id}")
+        
+        # Update state
+        self.agent_states[agent_id] = AgentState.WAKING
+        
+        metadata = self.hibernated_agents[agent_id]
+        
+        try:
+            # Load state from disk
+            if metadata.format == HibernationFormat.JSON:
+                with open(metadata.state_file_path, 'r') as f:
+                    state = json.load(f)
+            else:  # PICKLE format — now uses JSON internally; verify HMAC before deserializing
+                sig_path = metadata.state_file_path + ".sig"
+                if not os.path.exists(sig_path):
+                    raise ValueError(
+                        f"Missing HMAC signature for {metadata.state_file_path} — "
+                        "state file may have been tampered with"
+                    )
+                with open(metadata.state_file_path, 'rb') as f:
+                    raw = f.read()
+                with open(sig_path, 'r', encoding='utf-8') as f:
+                    expected_sig = f.read().strip()
+                actual_sig = hmac.new(self._hmac_key, raw, hashlib.sha256).hexdigest()
+                if not hmac.compare_digest(actual_sig, expected_sig):
+                    raise ValueError(
+                        f"HMAC verification failed for {metadata.state_file_path} — "
+                        "state file has been tampered with"
+                    )
+                state = json.loads(raw.decode('utf-8'))
+            
+            # Deserialize state
+            restored_state = self.deserialize_agent_state(state)
+            
+            # Update tracking
+            self.agent_activity[agent_id] = datetime.now()
+            self.agent_states[agent_id] = AgentState.ACTIVE
+            
+            # Remove from hibernated agents
+            del self.hibernated_agents[agent_id]
+            
+            self.logger.info(f"Agent {agent_id} woken successfully")
+            
+            return restored_state
+            
+        except Exception as e:
+            self.logger.error(f"Failed to wake agent {agent_id}: {e}")
+            self.agent_states[agent_id] = AgentState.HIBERNATED
+            raise
+    
+    def is_agent_hibernated(self, agent_id: str) -> bool:
+        """Check if an agent is currently hibernated"""
+        return agent_id in self.hibernated_agents
+    
+    def get_hibernated_agents(self) -> List[HibernatedAgentMetadata]:
+        """Get list of all hibernated agents"""
+        return list(self.hibernated_agents.values())
+    
+    def record_agent_activity(self, agent_id: str):
+        """Record activity for an agent (resets idle timer)"""
+        self.agent_activity[agent_id] = datetime.now()
+        if agent_id not in self.agent_states or self.agent_states[agent_id] == AgentState.HIBERNATED:
+            self.agent_states[agent_id] = AgentState.ACTIVE
+    
+    def get_idle_agents(self, min_idle_seconds: Optional[int] = None) -> List[str]:
+        """
+        Get list of agent IDs that have been idle for the specified duration.
+        
+        Args:
+            min_idle_seconds: Minimum idle time in seconds (uses config default if None)
+            
+        Returns:
+            List of agent IDs that are idle
+        """
+        idle_threshold = min_idle_seconds or self.config.idle_timeout_seconds
+        now = datetime.now()
+        idle_agents = []
+        
+        for agent_id, last_activity in self.agent_activity.items():
+            # Skip already hibernated agents
+            if self.is_agent_hibernated(agent_id):
+                continue
+            
+            idle_time = (now - last_activity).total_seconds()
+            if idle_time >= idle_threshold:
+                idle_agents.append(agent_id)
+        
+        return idle_agents
+    
+    def cleanup_old_hibernated_agents(self, max_age_days: Optional[int] = None):
+        """
+        Clean up hibernated agents older than the specified age.
+        
+        Args:
+            max_age_days: Maximum age in days (uses config default if None)
+        """
+        max_age = max_age_days or self.config.auto_cleanup_days
+        now = datetime.now()
+        cleanup_threshold = timedelta(days=max_age)
+        
+        agents_to_cleanup = []
+        for agent_id, metadata in self.hibernated_agents.items():
+            age = now - metadata.hibernated_at
+            if age > cleanup_threshold:
+                agents_to_cleanup.append(agent_id)
+        
+        for agent_id in agents_to_cleanup:
+            self._cleanup_hibernated_agent(agent_id)
+    
+    def _cleanup_hibernated_agent(self, agent_id: str):
+        """Remove hibernated agent and delete its state file"""
+        if agent_id not in self.hibernated_agents:
+            return
+        
+        metadata = self.hibernated_agents[agent_id]
+        
+        try:
+            # Delete state file
+            if os.path.exists(metadata.state_file_path):
+                os.remove(metadata.state_file_path)
+            
+            # Remove from tracking
+            del self.hibernated_agents[agent_id]
+            if agent_id in self.agent_states:
+                del self.agent_states[agent_id]
+            if agent_id in self.agent_activity:
+                del self.agent_activity[agent_id]
+            
+            self.logger.info(f"Cleaned up hibernated agent {agent_id}")
+            
+        except Exception as e:
+            self.logger.error(f"Failed to cleanup hibernated agent {agent_id}: {e}")
+    
+    def get_statistics(self) -> Dict[str, Any]:
+        """Get statistics about hibernation"""
+        total_hibernated = len(self.hibernated_agents)
+        total_size = sum(m.state_size_bytes for m in self.hibernated_agents.values())
+        
+        return {
+            "total_hibernated_agents": total_hibernated,
+            "total_state_size_bytes": total_size,
+            "total_state_size_mb": total_size / (1024 * 1024),
+            "active_agents": len([s for s in self.agent_states.values() if s == AgentState.ACTIVE]),
+            "storage_path": self.config.storage_path,
+            "format": self.config.format.value,
+            "idle_timeout_seconds": self.config.idle_timeout_seconds
+        }