claude-mpm 3.4.10__py3-none-any.whl → 5.4.55__py3-none-any.whl

claude_mpm/core/agent_session_manager.py

@@ -1,11 +1,11 @@
 """Agent-specific session management for performance optimization."""
 
-import uuid
-from typing import Dict, Optional, Tuple
-from datetime import datetime, timedelta
-from pathlib import Path
 import json
+import uuid
 from collections import defaultdict
+from datetime import datetime, timedelta, timezone
+from pathlib import Path
+from typing import Dict, Optional
 
 from ..core.logger import get_logger
 
@@ -14,10 +14,12 @@ logger = get_logger(__name__)
 
 class AgentSessionManager:
     """Manages separate sessions for each agent type to avoid context pollution."""
-    
-    def __init__(self, session_dir: Optional[Path] = None, max_sessions_per_agent: int = 3):
+
+    def __init__(
+        self, session_dir: Optional[Path] = None, max_sessions_per_agent: int = 3
+    ):
         """Initialize agent session manager.
-        
+
         Args:
             session_dir: Directory to store session metadata
             max_sessions_per_agent: Maximum concurrent sessions per agent type
@@ -25,129 +27,136 @@ class AgentSessionManager:
         self.session_dir = session_dir or Path.home() / ".claude-mpm" / "agent_sessions"
         self.session_dir.mkdir(parents=True, exist_ok=True)
         self.max_sessions_per_agent = max_sessions_per_agent
-        
+
         # Sessions organized by agent type
         self.agent_sessions: Dict[str, Dict[str, Dict]] = defaultdict(dict)
         self.session_locks: Dict[str, bool] = {}  # Track which sessions are in use
-        
+
         self._load_sessions()
-    
-    def get_agent_session(self, agent_type: str, create_if_missing: bool = True) -> Optional[str]:
+
+    def get_agent_session(
+        self, agent_type: str, create_if_missing: bool = True
+    ) -> Optional[str]:
         """Get or create a session for a specific agent type.
-        
+
         Args:
             agent_type: Type of agent (e.g., "engineer", "qa", "documentation")
             create_if_missing: Whether to create a new session if none available
-            
+
         Returns:
             Session ID or None
         """
         agent_type = agent_type.lower()
-        
+
         # Find an available session for this agent
         for session_id, session_data in self.agent_sessions[agent_type].items():
             if not self.session_locks.get(session_id, False):
                 # Check if session is still fresh (not too old)
                 created = datetime.fromisoformat(session_data["created_at"])
-                if datetime.now() - created < timedelta(hours=1):
+                if datetime.now(timezone.utc) - created < timedelta(hours=1):
                     # Use this session
                     self.session_locks[session_id] = True
-                    session_data["last_used"] = datetime.now().isoformat()
+                    session_data["last_used"] = datetime.now(timezone.utc).isoformat()
                     session_data["use_count"] += 1
                     logger.info(f"Reusing session {session_id} for {agent_type} agent")
                     return session_id
-        
+
         # No available session, create new one if allowed
-        if create_if_missing and len(self.agent_sessions[agent_type]) < self.max_sessions_per_agent:
+        if (
+            create_if_missing
+            and len(self.agent_sessions[agent_type]) < self.max_sessions_per_agent
+        ):
             return self.create_agent_session(agent_type)
-        
+
         return None
-    
+
     def create_agent_session(self, agent_type: str) -> str:
         """Create a new session for a specific agent type.
-        
+
         Args:
             agent_type: Type of agent
-            
+
         Returns:
             New session ID
         """
         session_id = str(uuid.uuid4())
         agent_type = agent_type.lower()
-        
+
         session_data = {
             "id": session_id,
             "agent_type": agent_type,
-            "created_at": datetime.now().isoformat(),
-            "last_used": datetime.now().isoformat(),
+            "created_at": datetime.now(timezone.utc).isoformat(),
+            "last_used": datetime.now(timezone.utc).isoformat(),
             "use_count": 0,
-            "tasks_completed": []
+            "tasks_completed": [],
         }
-        
+
         self.agent_sessions[agent_type][session_id] = session_data
         self.session_locks[session_id] = True
         self._save_sessions()
-        
+
         logger.info(f"Created new session {session_id} for {agent_type} agent")
         return session_id
-    
+
     def release_session(self, session_id: str):
         """Release a session back to the pool.
-        
+
         Args:
             session_id: Session to release
         """
         if session_id in self.session_locks:
             self.session_locks[session_id] = False
             logger.debug(f"Released session {session_id}")
-    
+
     def record_task(self, session_id: str, task: str, success: bool = True):
         """Record a task completion for a session.
-        
+
         Args:
             session_id: Session ID
             task: Task description
             success: Whether task completed successfully
         """
         # Find which agent this session belongs to
-        for agent_type, sessions in self.agent_sessions.items():
+        for _agent_type, sessions in self.agent_sessions.items():
             if session_id in sessions:
-                sessions[session_id]["tasks_completed"].append({
-                    "task": task[:100],  # Truncate long tasks
-                    "timestamp": datetime.now().isoformat(),
-                    "success": success
-                })
+                sessions[session_id]["tasks_completed"].append(
+                    {
+                        "task": task[:100],  # Truncate long tasks
+                        "timestamp": datetime.now(timezone.utc).isoformat(),
+                        "success": success,
+                    }
+                )
                 self._save_sessions()
                 break
-    
+
     def cleanup_old_sessions(self, max_age_hours: int = 4):
         """Remove sessions older than max_age_hours.
-        
+
         Args:
             max_age_hours: Maximum age in hours
         """
-        now = datetime.now()
+        now = datetime.now(timezone.utc)
         max_age = timedelta(hours=max_age_hours)
-        
+
         for agent_type in list(self.agent_sessions.keys()):
             expired = []
             for session_id, session_data in self.agent_sessions[agent_type].items():
                 created = datetime.fromisoformat(session_data["created_at"])
                 if now - created > max_age:
                     expired.append(session_id)
-            
+
             for session_id in expired:
                 del self.agent_sessions[agent_type][session_id]
                 if session_id in self.session_locks:
                     del self.session_locks[session_id]
                 logger.info(f"Cleaned up expired {agent_type} session: {session_id}")
-        
+
         if any(self.agent_sessions.values()):
             self._save_sessions()
-    
+
     def get_session_stats(self) -> Dict[str, Dict]:
         """Get statistics about current sessions.
-        
+
         Returns:
             Stats by agent type
         """
@@ -155,23 +164,23 @@ class AgentSessionManager:
         for agent_type, sessions in self.agent_sessions.items():
             active = sum(1 for sid in sessions if self.session_locks.get(sid, False))
             total_tasks = sum(len(s["tasks_completed"]) for s in sessions.values())
-            
+
             stats[agent_type] = {
                 "total_sessions": len(sessions),
                 "active_sessions": active,
                 "available_sessions": len(sessions) - active,
-                "total_tasks_completed": total_tasks
+                "total_tasks_completed": total_tasks,
             }
-        
+
         return stats
-    
+
     def initialize_agent_session(self, agent_type: str, session_id: str) -> str:
         """Initialize a session with agent-specific context.
-        
+
         Args:
             agent_type: Type of agent
             session_id: Session ID to initialize
-            
+
         Returns:
             Initialization prompt for the agent
         """
@@ -184,36 +193,38 @@ class AgentSessionManager:
             "security": "You are the Security Agent. Your role is security analysis and protection. Confirm with 'Security Agent ready.'",
             "ops": "You are the Ops Agent. Your role is deployment and operations. Confirm with 'Ops Agent ready.'",
             "data_engineer": "You are the Data Engineer Agent. Your role is data management and processing. Confirm with 'Data Engineer Agent ready.'",
-            "version_control": "You are the Version Control Agent. Your role is Git operations and version management. Confirm with 'Version Control Agent ready.'"
+            "version_control": "You are the Version Control Agent. Your role is Git operations and version management. Confirm with 'Version Control Agent ready.'",
         }
-        
+
         return initialization_prompts.get(
-            agent_type.lower(), 
-            f"You are the {agent_type} Agent. Confirm with '{agent_type} Agent ready.'"
+            agent_type.lower(),
+            f"You are the {agent_type} Agent. Confirm with '{agent_type} Agent ready.'",
         )
-    
+
     def _save_sessions(self):
         """Save sessions to disk."""
         session_file = self.session_dir / "agent_sessions.json"
         try:
             data = {
                 "agent_sessions": dict(self.agent_sessions),
-                "updated_at": datetime.now().isoformat()
+                "updated_at": datetime.now(timezone.utc).isoformat(),
             }
-            with open(session_file, 'w') as f:
+            with session_file.open("w") as f:
                 json.dump(data, f, indent=2)
         except Exception as e:
             logger.error(f"Failed to save agent sessions: {e}")
-    
+
     def _load_sessions(self):
         """Load sessions from disk."""
         session_file = self.session_dir / "agent_sessions.json"
         if session_file.exists():
             try:
-                with open(session_file, 'r') as f:
+                with session_file.open() as f:
                     data = json.load(f)
-                    self.agent_sessions = defaultdict(dict, data.get("agent_sessions", {}))
-                
+                    self.agent_sessions = defaultdict(
+                        dict, data.get("agent_sessions", {})
+                    )
+
                 # Clean up old sessions on load
                 self.cleanup_old_sessions()
             except Exception as e:
@@ -228,7 +239,7 @@ class SubprocessOrchestrator:
         self.agent_session_manager = AgentSessionManager()
         # Pre-warm common agent sessions on startup
         self._prewarm_agent_sessions()
-    
+
     def _prewarm_agent_sessions(self):
         '''Pre-warm sessions for common agents.'''
         common_agents = ['engineer', 'qa', 'documentation']
@@ -242,11 +253,11 @@ class SubprocessOrchestrator:
                 timeout=10
             )
             self.agent_session_manager.release_session(session_id)
-    
+
     def run_subprocess(self, agent: str, task: str) -> Tuple[str, float, int]:
         # Get a session for this agent type
         session_id = self.agent_session_manager.get_agent_session(agent)
-        
+
         try:
             # Create agent prompt (without role definition since session knows)
             prompt = f'''
@@ -256,18 +267,18 @@ class SubprocessOrchestrator:
 ## Response Format
 Provide a clear, structured response that completes the requested task.
 '''
-            
+
             # Run with agent-specific session
             stdout, stderr, returncode = self.launcher.launch_oneshot(
                 message=prompt,
                 session_id=session_id,  # Reuses agent-specific context!
                 timeout=60
             )
-            
+
             # Record task completion
             self.agent_session_manager.record_task(session_id, task, returncode == 0)
-            
+
         finally:
             # Release session back to pool
             self.agent_session_manager.release_session(session_id)
-"""
+"""

claude_mpm/core/api_validator.py

@@ -0,0 +1,330 @@
+"""API Key Validation Module for Claude MPM.
+
+This module validates API keys for various services on startup to ensure
+proper configuration and prevent runtime failures. It follows the principle
+of failing fast with clear error messages rather than degrading gracefully.
+"""
+
+import os
+from typing import Dict, List, Optional, Tuple
+
+import requests
+
+from claude_mpm.core.logger import get_logger
+
+
+class APIKeyValidator:
+    """Validates API keys for various services on framework startup."""
+
+    def __init__(self, config: Optional[Dict] = None):
+        """Initialize the API validator.
+
+        Args:
+            config: Optional configuration dictionary
+        """
+        self.logger = get_logger("api_validator")
+        self.config = config or {}
+        self.errors: List[str] = []
+        self.warnings: List[str] = []
+
+    def validate_all_keys(
+        self, strict: bool = True
+    ) -> Tuple[bool, List[str], List[str]]:
+        """Validate all configured API keys.
+
+        Args:
+            strict: If True, validation failures raise exceptions.
+                   If False, failures are logged as warnings.
+
+        Returns:
+            Tuple of (success, errors, warnings)
+        """
+        self.errors = []
+        self.warnings = []
+
+        # Check if validation is enabled
+        if not self.config.get("validate_api_keys", True):
+            self.logger.info("API key validation disabled in config")
+            return True, [], []
+
+        # Validate OpenAI key if configured
+        openai_key = os.getenv("OPENAI_API_KEY")
+        if openai_key:
+            self._validate_openai_key(openai_key)
+
+        # Validate Anthropic key if configured
+        anthropic_key = os.getenv("ANTHROPIC_API_KEY")
+        if anthropic_key:
+            self._validate_anthropic_key(anthropic_key)
+
+        # Validate GitHub token if configured
+        github_token = os.getenv("GITHUB_TOKEN")
+        if github_token:
+            self._validate_github_token(github_token)
+
+        # Validate custom API keys from config
+        custom_apis = self.config.get("custom_api_validations", {})
+        for api_name, validation_config in custom_apis.items():
+            self._validate_custom_api(api_name, validation_config)
+
+        # Report results
+        if self.errors:
+            error_msg = "API Key Validation Failed:\n" + "\n".join(self.errors)
+            if strict:
+                self.logger.error(error_msg)
+                raise ValueError(error_msg)
+            self.logger.warning(error_msg)
+
+        if self.warnings:
+            for warning in self.warnings:
+                self.logger.warning(warning)
+
+        if not self.errors:
+            self.logger.info("✅ All configured API keys validated successfully")
+
+        return not bool(self.errors), self.errors, self.warnings
+
+    def _validate_openai_key(self, api_key: str) -> bool:
+        """Validate OpenAI API key.
+
+        Args:
+            api_key: The OpenAI API key to validate
+
+        Returns:
+            True if valid, False otherwise
+        """
+        try:
+            # Make a lightweight request to validate the key
+            response = requests.get(
+                "https://api.openai.com/v1/models",
+                headers={"Authorization": f"Bearer {api_key}"},
+                timeout=10,
+            )
+
+            if response.status_code == 401:
+                self.errors.append("❌ OpenAI API key is invalid (401 Unauthorized)")
+                return False
+            if response.status_code == 403:
+                self.errors.append(
+                    "❌ OpenAI API key lacks required permissions (403 Forbidden)"
+                )
+                return False
+            if response.status_code == 429:
+                # Rate limited but key is valid
+                self.warnings.append("⚠️ OpenAI API key is valid but rate limited")
+                return True
+            if response.status_code == 200:
+                self.logger.debug("✅ OpenAI API key validated successfully")
+                return True
+            self.warnings.append(
+                f"⚠️ OpenAI API returned unexpected status: {response.status_code}"
+            )
+            return True  # Assume valid for unexpected status codes
+
+        except requests.exceptions.Timeout:
+            self.warnings.append(
+                "⚠️ OpenAI API validation timed out - assuming key is valid"
+            )
+            return True
+        except requests.exceptions.ConnectionError as e:
+            self.warnings.append(f"⚠️ Could not connect to OpenAI API: {e}")
+            return True
+        except Exception as e:
+            self.errors.append(f"❌ OpenAI API validation failed with error: {e}")
+            return False
+
+    def _validate_anthropic_key(self, api_key: str) -> bool:
+        """Validate Anthropic API key.
+
+        Args:
+            api_key: The Anthropic API key to validate
+
+        Returns:
+            True if valid, False otherwise
+        """
+        try:
+            # Make a minimal request to validate the key
+            # Using a very small max_tokens to minimize cost
+            response = requests.post(
+                "https://api.anthropic.com/v1/messages",
+                headers={
+                    "x-api-key": api_key,
+                    "anthropic-version": "2023-06-01",
+                    "content-type": "application/json",
+                },
+                json={
+                    "model": "claude-3-haiku-20240307",  # Use cheapest model
+                    "messages": [{"role": "user", "content": "test"}],
+                    "max_tokens": 1,
+                },
+                timeout=10,
+            )
+
+            if response.status_code == 401:
+                self.errors.append("❌ Anthropic API key is invalid (401 Unauthorized)")
+                return False
+            if response.status_code == 403:
+                self.errors.append(
+                    "❌ Anthropic API key lacks required permissions (403 Forbidden)"
+                )
+                return False
+            if response.status_code == 400:
+                # Bad request but key is valid (we sent minimal request on purpose)
+                self.logger.debug("✅ Anthropic API key validated successfully")
+                return True
+            if response.status_code == 429:
+                # Rate limited but key is valid
+                self.warnings.append("⚠️ Anthropic API key is valid but rate limited")
+                return True
+            if response.status_code == 200:
+                self.logger.debug("✅ Anthropic API key validated successfully")
+                return True
+            self.warnings.append(
+                f"⚠️ Anthropic API returned unexpected status: {response.status_code}"
+            )
+            return True
+
+        except requests.exceptions.Timeout:
+            self.warnings.append(
+                "⚠️ Anthropic API validation timed out - assuming key is valid"
+            )
+            return True
+        except requests.exceptions.ConnectionError as e:
+            self.warnings.append(f"⚠️ Could not connect to Anthropic API: {e}")
+            return True
+        except Exception as e:
+            self.errors.append(f"❌ Anthropic API validation failed with error: {e}")
+            return False
+
+    def _validate_github_token(self, token: str) -> bool:
+        """Validate GitHub personal access token.
+
+        Args:
+            token: The GitHub token to validate
+
+        Returns:
+            True if valid, False otherwise
+        """
+        try:
+            # Check token validity with minimal request
+            response = requests.get(
+                "https://api.github.com/user",
+                headers={
+                    "Authorization": f"token {token}",
+                    "Accept": "application/vnd.github.v3+json",
+                },
+                timeout=10,
+            )
+
+            if response.status_code == 401:
+                self.errors.append("❌ GitHub token is invalid (401 Unauthorized)")
+                return False
+            if response.status_code == 403:
+                self.errors.append(
+                    "❌ GitHub token lacks required permissions (403 Forbidden)"
+                )
+                return False
+            if response.status_code == 200:
+                self.logger.debug("✅ GitHub token validated successfully")
+                return True
+            self.warnings.append(
+                f"⚠️ GitHub API returned unexpected status: {response.status_code}"
+            )
+            return True
+
+        except requests.exceptions.Timeout:
+            self.warnings.append(
+                "⚠️ GitHub API validation timed out - assuming token is valid"
+            )
+            return True
+        except requests.exceptions.ConnectionError as e:
+            self.warnings.append(f"⚠️ Could not connect to GitHub API: {e}")
+            return True
+        except Exception as e:
+            self.errors.append(f"❌ GitHub token validation failed with error: {e}")
+            return False
+
+    def _validate_custom_api(self, api_name: str, validation_config: Dict) -> bool:
+        """Validate a custom API key based on configuration.
+
+        Args:
+            api_name: Name of the API
+            validation_config: Configuration for validating this API
+
+        Returns:
+            True if valid, False otherwise
+        """
+        try:
+            env_var = validation_config.get("env_var")
+            if not env_var:
+                return True
+
+            api_key = os.getenv(env_var)
+            if not api_key:
+                return True  # Not configured, skip validation
+
+            # Get validation endpoint and method
+            endpoint = validation_config.get("endpoint")
+            method = validation_config.get("method", "GET").upper()
+            headers = validation_config.get("headers", {})
+
+            # Replace {API_KEY} placeholder in headers
+            for key, value in headers.items():
+                if isinstance(value, str):
+                    headers[key] = value.replace("{API_KEY}", api_key)
+
+            # Make validation request
+            if method == "GET":
+                response = requests.get(endpoint, headers=headers, timeout=10)
+            elif method == "POST":
+                body = validation_config.get("body", {})
+                response = requests.post(
+                    endpoint, headers=headers, json=body, timeout=10
+                )
+            else:
+                self.warnings.append(
+                    f"⚠️ Unsupported validation method for {api_name}: {method}"
+                )
+                return True
+
+            # Check expected status codes
+            valid_status_codes = validation_config.get("valid_status_codes", [200])
+            if response.status_code in valid_status_codes:
+                self.logger.debug(f"✅ {api_name} API key validated successfully")
+                return True
+            if response.status_code == 401:
+                self.errors.append(
+                    f"❌ {api_name} API key is invalid (401 Unauthorized)"
+                )
+                return False
+            if response.status_code == 403:
+                self.errors.append(
+                    f"❌ {api_name} API key lacks permissions (403 Forbidden)"
+                )
+                return False
+            self.warnings.append(
+                f"⚠️ {api_name} API returned status: {response.status_code}"
+            )
+            return True
+
+        except Exception as e:
+            self.warnings.append(f"⚠️ {api_name} API validation failed: {e}")
+            return True
+
+
+def validate_api_keys(config: Optional[Dict] = None, strict: bool = True) -> bool:
+    """Convenience function to validate all API keys.
+
+    Args:
+        config: Optional configuration dictionary
+        strict: If True, raise exception on validation failure
+
+    Returns:
+        True if all validations passed, False otherwise
+
+    Raises:
+        ValueError: If strict=True and any validation fails
+    """
+    validator = APIKeyValidator(config)
+    success, _errors, _warnings = validator.validate_all_keys(strict=strict)
+    return success