hanzo-mcp 0.6.13__py3-none-any.whl → 0.7.1__py3-none-any.whl

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

hanzo_mcp/tools/agent/code_auth_tool.py

@@ -0,0 +1,194 @@
+"""Claude Code authentication tool.
+
+This tool manages API keys and accounts for Claude Code and other AI coding tools.
+"""
+
+from typing import Annotated, Optional, TypedDict, Unpack, final, override
+from mcp.server import FastMCP
+from mcp.server.fastmcp import Context as MCPContext
+from pydantic import Field
+
+from hanzo_mcp.tools.common.base import BaseTool
+from hanzo_mcp.tools.common.context import create_tool_context
+from hanzo_mcp.tools.agent.code_auth import CodeAuthManager
+
+
+class CodeAuthParams(TypedDict, total=False):
+    """Parameters for code auth tool."""
+    action: str
+    account: Optional[str]
+    provider: Optional[str]
+    api_key: Optional[str]
+    model: Optional[str]
+    description: Optional[str]
+    agent_id: Optional[str]
+    parent_account: Optional[str]
+
+
+@final
+class CodeAuthTool(BaseTool):
+    """Tool for managing Claude Code authentication and API keys."""
+    
+    @property
+    @override
+    def name(self) -> str:
+        """Get the tool name."""
+        return "code_auth"
+    
+    @property
+    @override
+    def description(self) -> str:
+        """Get the tool description."""
+        return """Manage Claude Code and AI provider authentication.
+
+Actions:
+- status: Show current login status
+- list: List all accounts
+- create: Create a new account
+- login: Login to an account
+- logout: Logout current account
+- switch: Switch between accounts
+- agent: Create/get agent account
+
+Examples:
+code_auth status
+code_auth list
+code_auth create --account work --provider claude
+code_auth login --account work
+code_auth logout
+code_auth switch --account personal
+code_auth agent --agent_id swarm_1 --parent_account work
+
+Providers: claude, openai, azure, deepseek, google, groq"""
+    
+    def __init__(self):
+        """Initialize the code auth tool."""
+        self.auth_manager = CodeAuthManager()
+    
+    @override
+    async def call(
+        self,
+        ctx: MCPContext,
+        **params: Unpack[CodeAuthParams],
+    ) -> str:
+        """Execute the code auth tool.
+        
+        Args:
+            ctx: MCP context
+            **params: Tool parameters
+            
+        Returns:
+            Result message
+        """
+        tool_ctx = create_tool_context(ctx)
+        await tool_ctx.set_tool_info(self.name)
+        
+        action = params.get("action", "status")
+        
+        if action == "status":
+            current = self.auth_manager.get_active_account()
+            if current:
+                info = self.auth_manager.get_account_info(current)
+                if info:
+                    return f"Logged in as: {current} ({info['provider']})"
+            return "Not logged in"
+        
+        elif action == "list":
+            accounts = self.auth_manager.list_accounts()
+            if not accounts:
+                return "No accounts configured"
+            
+            current = self.auth_manager.get_active_account()
+            lines = ["Configured accounts:"]
+            for account in accounts:
+                info = self.auth_manager.get_account_info(account)
+                marker = " (active)" if account == current else ""
+                lines.append(f"  - {account}: {info['provider']}{marker}")
+            return "\n".join(lines)
+        
+        elif action == "create":
+            account = params.get("account")
+            if not account:
+                return "Error: account name required"
+            
+            provider = params.get("provider", "claude")
+            api_key = params.get("api_key")
+            model = params.get("model")
+            description = params.get("description")
+            
+            success, msg = self.auth_manager.create_account(
+                account, provider, api_key, model, description
+            )
+            return msg
+        
+        elif action == "login":
+            account = params.get("account", "default")
+            success, msg = self.auth_manager.login(account)
+            return msg
+        
+        elif action == "logout":
+            success, msg = self.auth_manager.logout()
+            return msg
+        
+        elif action == "switch":
+            account = params.get("account")
+            if not account:
+                return "Error: account name required"
+            
+            success, msg = self.auth_manager.switch_account(account)
+            return msg
+        
+        elif action == "agent":
+            agent_id = params.get("agent_id")
+            if not agent_id:
+                return "Error: agent_id required"
+            
+            provider = params.get("provider", "claude")
+            parent_account = params.get("parent_account")
+            
+            # Try to create agent account
+            success, result = self.auth_manager.create_agent_account(
+                agent_id, provider, parent_account
+            )
+            
+            if success:
+                # Get credentials
+                creds = self.auth_manager.get_agent_credentials(agent_id)
+                if creds:
+                    return f"Agent account ready: {result} ({creds.provider})"
+                else:
+                    return f"Agent account created but no credentials: {result}"
+            else:
+                return f"Failed to create agent account: {result}"
+        
+        else:
+            return f"Unknown action: {action}. Use: status, list, create, login, logout, switch, agent"
+    
+    @override
+    def register(self, mcp_server: FastMCP) -> None:
+        """Register this tool with the MCP server."""
+        tool_self = self
+        
+        @mcp_server.tool(name=self.name, description=self.description)
+        async def code_auth(
+            ctx: MCPContext,
+            action: str = "status",
+            account: Optional[str] = None,
+            provider: Optional[str] = None,
+            api_key: Optional[str] = None,
+            model: Optional[str] = None,
+            description: Optional[str] = None,
+            agent_id: Optional[str] = None,
+            parent_account: Optional[str] = None,
+        ) -> str:
+            return await tool_self.call(
+                ctx,
+                action=action,
+                account=account,
+                provider=provider,
+                api_key=api_key,
+                model=model,
+                description=description,
+                agent_id=agent_id,
+                parent_account=parent_account,
+            )