hanzo-mcp 0.6.12__py3-none-any.whl → 0.7.0__py3-none-any.whl

hanzo_mcp/tools/agent/cli_agent_base.py

@@ -0,0 +1,191 @@
+"""Base class for CLI-based AI agent tools.
+
+This provides common functionality for spawning CLI-based AI coding assistants
+like Claude Code, OpenAI Codex, Google Gemini, and Grok.
+"""
+
+import asyncio
+import json
+import os
+import shutil
+import subprocess
+import tempfile
+from abc import abstractmethod
+from pathlib import Path
+from typing import Optional, Dict, Any, List, Tuple
+from mcp.server.fastmcp import Context as MCPContext
+
+from hanzo_mcp.tools.common.base import BaseTool
+from hanzo_mcp.tools.common.context import create_tool_context
+from hanzo_mcp.tools.common.permissions import PermissionManager
+
+
+class CLIAgentBase(BaseTool):
+    """Base class for CLI-based AI agent tools."""
+    
+    def __init__(
+        self,
+        permission_manager: PermissionManager,
+        command_name: str,
+        provider_name: str,
+        default_model: Optional[str] = None,
+        env_vars: Optional[List[str]] = None,
+        **kwargs
+    ):
+        """Initialize CLI agent base.
+        
+        Args:
+            permission_manager: Permission manager for access control
+            command_name: The CLI command name (e.g., 'claude', 'openai')
+            provider_name: The provider name (e.g., 'Claude', 'OpenAI')
+            default_model: Default model to use
+            env_vars: List of environment variables to check for API keys
+            **kwargs: Additional arguments
+        """
+        self.permission_manager = permission_manager
+        self.command_name = command_name
+        self.provider_name = provider_name
+        self.default_model = default_model
+        self.env_vars = env_vars or []
+        
+    def is_installed(self) -> bool:
+        """Check if the CLI tool is installed."""
+        return shutil.which(self.command_name) is not None
+    
+    def has_api_key(self) -> bool:
+        """Check if API key is available in environment."""
+        if not self.env_vars:
+            return True  # No API key needed
+        
+        for var in self.env_vars:
+            if os.environ.get(var):
+                return True
+        return False
+    
+    @abstractmethod
+    def get_cli_args(self, prompt: str, **kwargs) -> List[str]:
+        """Get CLI arguments for the specific tool.
+        
+        Args:
+            prompt: The prompt to send
+            **kwargs: Additional arguments
+            
+        Returns:
+            List of command arguments
+        """
+        pass
+    
+    async def execute_cli(
+        self,
+        ctx: MCPContext,
+        prompt: str,
+        working_dir: Optional[str] = None,
+        timeout: int = 300,
+        **kwargs
+    ) -> str:
+        """Execute the CLI command.
+        
+        Args:
+            ctx: MCP context
+            prompt: The prompt to send
+            working_dir: Working directory for the command
+            timeout: Command timeout in seconds
+            **kwargs: Additional arguments
+            
+        Returns:
+            Command output
+        """
+        tool_ctx = create_tool_context(ctx)
+        
+        # Check if installed
+        if not self.is_installed():
+            error_msg = f"{self.provider_name} CLI ({self.command_name}) is not installed. "
+            error_msg += f"Please install it first: https://github.com/anthropics/{self.command_name}"
+            await tool_ctx.error(error_msg)
+            return f"Error: {error_msg}"
+        
+        # Check API key if needed
+        if not self.has_api_key():
+            error_msg = f"No API key found for {self.provider_name}. "
+            error_msg += f"Set one of: {', '.join(self.env_vars)}"
+            await tool_ctx.error(error_msg)
+            return f"Error: {error_msg}"
+        
+        # Get command arguments
+        cli_args = self.get_cli_args(prompt, **kwargs)
+        
+        # Log command
+        await tool_ctx.info(f"Executing {self.provider_name}: {self.command_name} {' '.join(cli_args[:3])}...")
+        
+        try:
+            # Create temp file for prompt if needed
+            with tempfile.NamedTemporaryFile(mode='w', suffix='.txt', delete=False) as f:
+                f.write(prompt)
+                prompt_file = f.name
+            
+            # Some CLIs might need the prompt via stdin or file
+            if '--prompt-file' in cli_args:
+                # Replace placeholder with actual file
+                cli_args = [arg.replace('--prompt-file', prompt_file) if arg == '--prompt-file' else arg for arg in cli_args]
+            
+            # Execute command
+            process = await asyncio.create_subprocess_exec(
+                self.command_name,
+                *cli_args,
+                stdout=asyncio.subprocess.PIPE,
+                stderr=asyncio.subprocess.PIPE,
+                stdin=asyncio.subprocess.PIPE,
+                cwd=working_dir or os.getcwd()
+            )
+            
+            # Send prompt via stdin if not using file
+            if '--prompt-file' not in cli_args:
+                stdout, stderr = await asyncio.wait_for(
+                    process.communicate(input=prompt.encode()),
+                    timeout=timeout
+                )
+            else:
+                stdout, stderr = await asyncio.wait_for(
+                    process.communicate(),
+                    timeout=timeout
+                )
+            
+            # Clean up temp file
+            try:
+                os.unlink(prompt_file)
+            except:
+                pass
+            
+            if process.returncode != 0:
+                error_msg = stderr.decode() if stderr else "Unknown error"
+                await tool_ctx.error(f"{self.provider_name} failed: {error_msg}")
+                return f"Error: {error_msg}"
+            
+            result = stdout.decode()
+            await tool_ctx.info(f"{self.provider_name} completed successfully")
+            return result
+            
+        except asyncio.TimeoutError:
+            await tool_ctx.error(f"{self.provider_name} timed out after {timeout} seconds")
+            return f"Error: Command timed out after {timeout} seconds"
+        except Exception as e:
+            await tool_ctx.error(f"{self.provider_name} error: {str(e)}")
+            return f"Error: {str(e)}"
+    
+    async def call(
+        self,
+        ctx: MCPContext,
+        prompts: str,
+        **kwargs
+    ) -> str:
+        """Execute the CLI agent.
+        
+        Args:
+            ctx: MCP context
+            prompts: The prompt(s) to send
+            **kwargs: Additional arguments
+            
+        Returns:
+            Agent response
+        """
+        return await self.execute_cli(ctx, prompts, **kwargs)

hanzo_mcp/tools/agent/code_auth.py

@@ -0,0 +1,436 @@
+"""Claude Code and OpenAI Codex authentication management.
+
+This module provides tools to manage API keys and authentication for
+Claude Code CLI and OpenAI Codex, allowing separate accounts for swarm agents.
+"""
+
+import asyncio
+import os
+import subprocess
+import json
+import tempfile
+from pathlib import Path
+from typing import Optional, Tuple, Dict, Any, List
+from dataclasses import dataclass
+import keyring
+import getpass
+
+
+@dataclass
+class APICredential:
+    """API credential information."""
+    provider: str
+    api_key: str
+    model: Optional[str] = None
+    base_url: Optional[str] = None
+    org_id: Optional[str] = None
+    description: Optional[str] = None
+
+
+class CodeAuthManager:
+    """Manages authentication for Claude Code and other AI coding tools."""
+    
+    # Configuration paths
+    CONFIG_DIR = Path.home() / ".hanzo" / "auth"
+    ACCOUNTS_FILE = CONFIG_DIR / "accounts.json"
+    ACTIVE_ACCOUNT_FILE = CONFIG_DIR / "active_account"
+    
+    # Environment variable mappings
+    ENV_VARS = {
+        "claude": ["ANTHROPIC_API_KEY", "CLAUDE_API_KEY"],
+        "openai": ["OPENAI_API_KEY"],
+        "azure": ["AZURE_OPENAI_API_KEY", "AZURE_API_KEY"],
+        "deepseek": ["DEEPSEEK_API_KEY"],
+        "google": ["GOOGLE_API_KEY", "GEMINI_API_KEY"],
+        "groq": ["GROQ_API_KEY"],
+    }
+    
+    # Default models
+    DEFAULT_MODELS = {
+        "claude": "claude-3-5-sonnet-20241022",  # Latest Sonnet
+        "openai": "gpt-4o",
+        "azure": "gpt-4",
+        "deepseek": "deepseek-coder",
+        "google": "gemini-1.5-pro",
+        "groq": "llama3-70b-8192",
+    }
+    
+    def __init__(self):
+        """Initialize auth manager."""
+        self.ensure_config_dir()
+        self._env_backup = {}
+    
+    def ensure_config_dir(self):
+        """Ensure config directory exists."""
+        self.CONFIG_DIR.mkdir(parents=True, exist_ok=True)
+    
+    def get_active_account(self) -> Optional[str]:
+        """Get the currently active account."""
+        if self.ACTIVE_ACCOUNT_FILE.exists():
+            return self.ACTIVE_ACCOUNT_FILE.read_text().strip()
+        return "default"
+    
+    def set_active_account(self, account: str):
+        """Set the active account."""
+        self.ACTIVE_ACCOUNT_FILE.write_text(account)
+    
+    def _load_accounts(self) -> Dict[str, Dict[str, Any]]:
+        """Load all accounts."""
+        if not self.ACCOUNTS_FILE.exists():
+            return {}
+        
+        try:
+            with open(self.ACCOUNTS_FILE, 'r') as f:
+                return json.load(f)
+        except:
+            return {}
+    
+    def _save_accounts(self, accounts: Dict[str, Dict[str, Any]]):
+        """Save accounts."""
+        with open(self.ACCOUNTS_FILE, 'w') as f:
+            json.dump(accounts, f, indent=2)
+    
+    def list_accounts(self) -> List[str]:
+        """List all available accounts."""
+        accounts = self._load_accounts()
+        return list(accounts.keys())
+    
+    def get_account_info(self, account: str) -> Optional[Dict[str, Any]]:
+        """Get information about an account."""
+        accounts = self._load_accounts()
+        return accounts.get(account)
+    
+    def create_account(
+        self, 
+        account: str,
+        provider: str = "claude",
+        api_key: Optional[str] = None,
+        model: Optional[str] = None,
+        description: Optional[str] = None
+    ) -> Tuple[bool, str]:
+        """Create a new account.
+        
+        Args:
+            account: Account name
+            provider: Provider (claude, openai, etc.)
+            api_key: API key (will prompt if not provided)
+            model: Model to use (defaults to provider default)
+            description: Account description
+            
+        Returns:
+            Tuple of (success, message)
+        """
+        accounts = self._load_accounts()
+        
+        if account in accounts:
+            return False, f"Account '{account}' already exists"
+        
+        # Get API key if not provided
+        if not api_key:
+            api_key = self._prompt_for_api_key(provider)
+            if not api_key:
+                return False, "No API key provided"
+        
+        # Use default model if not specified
+        if not model:
+            model = self.DEFAULT_MODELS.get(provider)
+        
+        # Store in keyring for security
+        try:
+            keyring.set_password(f"hanzo-{provider}", account, api_key)
+        except:
+            # Fallback to file storage (less secure)
+            pass
+        
+        # Save account info
+        accounts[account] = {
+            "provider": provider,
+            "model": model,
+            "description": description or f"{provider} account",
+            "created_at": os.path.getmtime(__file__),
+            "has_keyring": self._has_keyring_support()
+        }
+        
+        self._save_accounts(accounts)
+        return True, f"Created account '{account}' for {provider}"
+    
+    def _prompt_for_api_key(self, provider: str) -> Optional[str]:
+        """Prompt user for API key."""
+        prompt = f"Enter {provider.upper()} API key: "
+        try:
+            return getpass.getpass(prompt)
+        except KeyboardInterrupt:
+            return None
+    
+    def _has_keyring_support(self) -> bool:
+        """Check if keyring is available."""
+        try:
+            keyring.get_keyring()
+            return True
+        except:
+            return False
+    
+    def login(self, account: str = "default") -> Tuple[bool, str]:
+        """Login to an account by setting environment variables.
+        
+        Args:
+            account: Account name to login to
+            
+        Returns:
+            Tuple of (success, message)
+        """
+        accounts = self._load_accounts()
+        
+        if account not in accounts:
+            return False, f"Account '{account}' not found"
+        
+        account_info = accounts[account]
+        provider = account_info["provider"]
+        
+        # Get API key from keyring or prompt
+        api_key = None
+        if account_info.get("has_keyring"):
+            try:
+                api_key = keyring.get_password(f"hanzo-{provider}", account)
+            except:
+                pass
+        
+        if not api_key:
+            # Try environment variable
+            for env_var in self.ENV_VARS.get(provider, []):
+                if env_var in os.environ:
+                    api_key = os.environ[env_var]
+                    break
+        
+        if not api_key:
+            api_key = self._prompt_for_api_key(provider)
+            if not api_key:
+                return False, "No API key available"
+        
+        # Backup current environment
+        self._backup_environment(provider)
+        
+        # Set environment variables
+        for env_var in self.ENV_VARS.get(provider, []):
+            os.environ[env_var] = api_key
+        
+        # Set active account
+        self.set_active_account(account)
+        
+        # Update shell if using claude command
+        self._update_claude_command(account_info)
+        
+        return True, f"Logged in as '{account}' ({provider})"
+    
+    def logout(self) -> Tuple[bool, str]:
+        """Logout by clearing environment variables."""
+        current = self.get_active_account()
+        
+        if not current or current == "default":
+            return False, "No active session"
+        
+        accounts = self._load_accounts()
+        if current not in accounts:
+            return False, f"Unknown account: {current}"
+        
+        provider = accounts[current]["provider"]
+        
+        # Clear environment variables
+        for env_var in self.ENV_VARS.get(provider, []):
+            if env_var in os.environ:
+                del os.environ[env_var]
+        
+        # Restore backed up environment if any
+        self._restore_environment(provider)
+        
+        # Clear active account
+        if self.ACTIVE_ACCOUNT_FILE.exists():
+            self.ACTIVE_ACCOUNT_FILE.unlink()
+        
+        return True, f"Logged out from '{current}'"
+    
+    def _backup_environment(self, provider: str):
+        """Backup current environment variables."""
+        for env_var in self.ENV_VARS.get(provider, []):
+            if env_var in os.environ:
+                self._env_backup[env_var] = os.environ[env_var]
+    
+    def _restore_environment(self, provider: str):
+        """Restore backed up environment variables."""
+        for env_var in self.ENV_VARS.get(provider, []):
+            if env_var in self._env_backup:
+                os.environ[env_var] = self._env_backup[env_var]
+                del self._env_backup[env_var]
+    
+    def _update_claude_command(self, account_info: Dict[str, Any]):
+        """Update claude command configuration if needed."""
+        # Check if claude command exists
+        try:
+            result = subprocess.run(
+                ["which", "claude"],
+                capture_output=True,
+                text=True
+            )
+            if result.returncode == 0:
+                # Claude command exists, update its config
+                claude_config = Path.home() / ".claude" / "config.json"
+                if claude_config.exists():
+                    try:
+                        with open(claude_config, 'r') as f:
+                            config = json.load(f)
+                        
+                        # Update model if specified
+                        if account_info.get("model"):
+                            config["default_model"] = account_info["model"]
+                        
+                        with open(claude_config, 'w') as f:
+                            json.dump(config, f, indent=2)
+                    except:
+                        pass
+        except:
+            pass
+    
+    def switch_account(self, account: str) -> Tuple[bool, str]:
+        """Switch to a different account."""
+        # Logout current
+        self.logout()
+        
+        # Login to new account
+        return self.login(account)
+    
+    def create_agent_account(
+        self, 
+        agent_id: str,
+        provider: str = "claude",
+        parent_account: Optional[str] = None
+    ) -> Tuple[bool, str]:
+        """Create an account for a swarm agent.
+        
+        Args:
+            agent_id: Unique agent identifier
+            provider: AI provider
+            parent_account: Parent account to clone from
+            
+        Returns:
+            Tuple of (success, account_name)
+        """
+        agent_account = f"agent_{agent_id}"
+        
+        # If parent account specified, clone its credentials
+        if parent_account:
+            parent_info = self.get_account_info(parent_account)
+            if not parent_info:
+                return False, f"Parent account '{parent_account}' not found"
+            
+            # Get parent API key
+            api_key = None
+            if parent_info.get("has_keyring"):
+                try:
+                    api_key = keyring.get_password(
+                        f"hanzo-{parent_info['provider']}", 
+                        parent_account
+                    )
+                except:
+                    pass
+            
+            if api_key:
+                success, msg = self.create_account(
+                    agent_account,
+                    provider=parent_info["provider"],
+                    api_key=api_key,
+                    model=parent_info.get("model"),
+                    description=f"Agent account (parent: {parent_account})"
+                )
+                if success:
+                    return True, agent_account
+        
+        # Create with current environment
+        for env_var in self.ENV_VARS.get(provider, []):
+            if env_var in os.environ:
+                success, msg = self.create_account(
+                    agent_account,
+                    provider=provider,
+                    api_key=os.environ[env_var],
+                    model=self.DEFAULT_MODELS.get(provider),
+                    description=f"Agent account for {agent_id}"
+                )
+                if success:
+                    return True, agent_account
+        
+        return False, "No credentials available for agent"
+    
+    def get_agent_credentials(self, agent_id: str) -> Optional[APICredential]:
+        """Get credentials for an agent.
+        
+        Args:
+            agent_id: Agent identifier
+            
+        Returns:
+            APICredential if found
+        """
+        agent_account = f"agent_{agent_id}"
+        account_info = self.get_account_info(agent_account)
+        
+        if not account_info:
+            return None
+        
+        # Get API key
+        api_key = None
+        provider = account_info["provider"]
+        
+        if account_info.get("has_keyring"):
+            try:
+                api_key = keyring.get_password(f"hanzo-{provider}", agent_account)
+            except:
+                pass
+        
+        if not api_key:
+            # Try current environment
+            for env_var in self.ENV_VARS.get(provider, []):
+                if env_var in os.environ:
+                    api_key = os.environ[env_var]
+                    break
+        
+        if not api_key:
+            return None
+        
+        return APICredential(
+            provider=provider,
+            api_key=api_key,
+            model=account_info.get("model"),
+            description=account_info.get("description")
+        )
+
+
+# Update swarm tool to use latest Sonnet
+def get_latest_claude_model() -> str:
+    """Get the latest Claude model identifier."""
+    # As of the knowledge cutoff, this is the latest Sonnet
+    # In production, this could query an API for the latest model
+    return "claude-3-5-sonnet-20241022"
+
+
+# Token counting using tiktoken (same as current implementation)
+def count_tokens_streaming(text_stream) -> int:
+    """Count tokens in a streaming fashion.
+    
+    This uses the same tiktoken approach as the truncate module,
+    but processes text as it streams.
+    """
+    import tiktoken
+    
+    try:
+        # Use cl100k_base encoding (Claude/GPT-4 compatible)
+        encoding = tiktoken.get_encoding("cl100k_base")
+    except:
+        # Fallback to simple estimation
+        return len(text_stream) // 4
+    
+    total_tokens = 0
+    for chunk in text_stream:
+        if isinstance(chunk, str):
+            total_tokens += len(encoding.encode(chunk))
+        elif isinstance(chunk, bytes):
+            total_tokens += len(encoding.encode(chunk.decode('utf-8', errors='ignore')))
+    
+    return total_tokens