hanzo-mcp 0.7.7__py3-none-any.whl → 0.8.1__py3-none-any.whl

hanzo_mcp/tools/agent/code_auth.py

@@ -4,21 +4,21 @@ 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
+import getpass
+import subprocess
+from typing import Any, Dict, List, Tuple, Optional
 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
@@ -29,12 +29,12 @@ class APICredential:
 
 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"],
@@ -44,7 +44,7 @@ class CodeAuthManager:
         "google": ["GOOGLE_API_KEY", "GEMINI_API_KEY"],
         "groq": ["GROQ_API_KEY"],
     }
-    
+
     # Default models
     DEFAULT_MODELS = {
         "claude": "claude-3-5-sonnet-20241022",  # Latest Sonnet
@@ -54,106 +54,106 @@ class CodeAuthManager:
         "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:
+            with open(self.ACCOUNTS_FILE, "r") as f:
                 return json.load(f)
-        except:
+        except Exception:
             return {}
-    
+
     def _save_accounts(self, accounts: Dict[str, Dict[str, Any]]):
         """Save accounts."""
-        with open(self.ACCOUNTS_FILE, 'w') as f:
+        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, 
+        self,
         account: str,
         provider: str = "claude",
         api_key: Optional[str] = None,
         model: Optional[str] = None,
-        description: 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:
+        except Exception:
             # 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()
+            "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: "
@@ -161,189 +161,184 @@ class CodeAuthManager:
             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:
+        except Exception:
             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:
+            except Exception:
                 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
-            )
+            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:
+                        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:
+
+                        with open(claude_config, "w") as f:
                             json.dump(config, f, indent=2)
-                    except:
+                    except Exception:
                         pass
-        except:
+        except Exception:
             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, 
+        self,
         agent_id: str,
         provider: str = "claude",
-        parent_account: Optional[str] = None
+        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
+                        f"hanzo-{parent_info['provider']}", parent_account
                     )
-                except:
+                except Exception:
                     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})"
+                    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:
@@ -352,53 +347,53 @@ class CodeAuthManager:
                     provider=provider,
                     api_key=os.environ[env_var],
                     model=self.DEFAULT_MODELS.get(provider),
-                    description=f"Agent account for {agent_id}"
+                    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:
+            except Exception:
                 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")
+            description=account_info.get("description"),
         )
 
 
@@ -413,24 +408,24 @@ def get_latest_claude_model() -> str:
 # 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:
+    except Exception:
         # 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
+            total_tokens += len(encoding.encode(chunk.decode("utf-8", errors="ignore")))
+
+    return total_tokens