alphai 0.1.2__py3-none-any.whl → 0.2.1__py3-none-any.whl

alphai/__init__.py

@@ -1,10 +1,48 @@
 """alphai - A CLI tool and Python package for the runalph.ai platform."""
 
-__version__ = "0.1.2"
+__version__ = "0.2.1"
 __author__ = "American Data Science"
 __email__ = "support@americandatascience.com"
 
 from .client import AlphAIClient
 from .config import Config
+from .jupyter_manager import JupyterManager
+from .exceptions import (
+    AlphAIException,
+    AuthenticationError,
+    AuthorizationError,
+    APIError,
+    DockerError,
+    DockerNotAvailableError,
+    ContainerError,
+    TunnelError,
+    CloudflaredError,
+    ConfigurationError,
+    ValidationError,
+    NetworkError,
+    TimeoutError,
+    ResourceNotFoundError,
+    JupyterError,
+)
 
-__all__ = ["AlphAIClient", "Config", "__version__"]
+__all__ = [
+    "AlphAIClient",
+    "Config",
+    "JupyterManager",
+    "__version__",
+    "AlphAIException",
+    "AuthenticationError",
+    "AuthorizationError",
+    "APIError",
+    "DockerError",
+    "DockerNotAvailableError",
+    "ContainerError",
+    "TunnelError",
+    "CloudflaredError",
+    "ConfigurationError",
+    "ValidationError",
+    "NetworkError",
+    "TimeoutError",
+    "ResourceNotFoundError",
+    "JupyterError",
+]

alphai/auth.py

@@ -16,6 +16,10 @@ from rich.prompt import Prompt
 from rich.panel import Panel
 
 from .config import Config
+from .utils import get_logger
+from . import exceptions
+
+logger = get_logger(__name__)
 
 
 class CallbackHandler(BaseHTTPRequestHandler):
@@ -47,7 +51,7 @@ class CallbackHandler(BaseHTTPRequestHandler):
                     else:
                         # No success page, redirect directly to dashboard
                         redirect_url = api_url
-            except:
+            except Exception:
                 # Network error or timeout, redirect directly to dashboard
                 redirect_url = api_url
             
@@ -99,15 +103,19 @@ class AuthManager:
     
     def login_with_token(self, token: str) -> bool:
         """Login with a provided bearer token."""
+        logger.info("Attempting login with provided token")
         if not token.strip():
+            logger.error("Empty token provided")
             self.console.print("[red]Error: Empty token provided[/red]")
-            return False
+            raise exceptions.ValidationError("Token cannot be empty")
         
         # Validate the token
         if self.validate_token(token):
             self.config.set_bearer_token(token)
+            logger.info("Login successful with token")
             return True
         else:
+            logger.error("Invalid token provided")
             self.console.print("[red]Error: Invalid token[/red]")
             return False
     
@@ -171,6 +179,7 @@ class AuthManager:
     
     def browser_login(self, port: int = 8080) -> bool:
         """Perform browser-based login using OAuth flow."""
+        logger.info("Starting browser-based authentication")
         # Try to find an available port
         for attempt_port in range(port, port + 10):
             try:
@@ -179,18 +188,21 @@ class AuthManager:
                 httpd.timeout = 60  # 1 minute timeout
                 httpd.token = None
                 httpd.error = None
-                httpd.api_url = self.config.api_url  # Pass api_url to server
+                httpd.api_url = self.config.base_url  # Pass base_url for redirects (not /api)
+                logger.debug(f"Callback server started on port {attempt_port}")
                 break
             except OSError:
+                logger.debug(f"Port {attempt_port} not available, trying next")
                 continue
         else:
+            logger.error("Could not find available port for OAuth callback")
             self.console.print("[red]Error: Could not find an available port for callback[/red]")
             return False
         
-        redirect_uri = f"http://localhost:{attempt_port}"
+        redirect_uri = f"http://127.0.0.1:{attempt_port}"
         
-        # Construct the authentication URL
-        auth_url = f"{self.config.api_url}/auth/cli"
+        # Construct the authentication URL (auth is at base URL, not /api)
+        auth_url = f"{self.config.base_url}/auth/cli"
         auth_params = {
             "redirect_uri": redirect_uri,
             "response_type": "token",
@@ -242,6 +254,7 @@ class AuthManager:
     
     def validate_token(self, token: str) -> bool:
         """Validate a bearer token by making a test API call."""
+        logger.debug("Validating bearer token")
         try:
             with httpx.Client() as client:
                 headers = {
@@ -251,31 +264,38 @@ class AuthManager:
                 
                 # Test the token by trying to get organizations
                 response = client.get(
-                    f"{self.config.api_url}/api/orgs",
+                    f"{self.config.api_url}/orgs",
                     headers=headers,
                     timeout=10.0
                 )
                 
                 # Check if the response is successful
                 if response.status_code in (200, 201):
+                    logger.info("Token validation successful")
                     return True
                 elif response.status_code == 401:
+                    logger.error("Token validation failed: Invalid or expired token")
                     self.console.print("[red]Error: Invalid or expired token[/red]")
                     return False
                 elif response.status_code == 403:
+                    logger.error("Token validation failed: Access forbidden")
                     self.console.print("[red]Error: Access forbidden - check your permissions[/red]")
                     return False
                 else:
+                    logger.error(f"Token validation failed: API returned status {response.status_code}")
                     self.console.print(f"[red]Error: API returned status {response.status_code}[/red]")
                     return False
                     
-        except httpx.ConnectError:
+        except httpx.ConnectError as e:
+            logger.error(f"Connection error during token validation: {e}")
             self.console.print(f"[red]Error: Could not connect to API at {self.config.api_url}[/red]")
             return False
-        except httpx.TimeoutException:
+        except httpx.TimeoutException as e:
+            logger.error(f"Timeout during token validation: {e}")
             self.console.print("[red]Error: Request timed out[/red]")
             return False
         except Exception as e:
+            logger.error(f"Unexpected error validating token: {e}", exc_info=True)
             self.console.print(f"[red]Error validating token: {e}[/red]")
             return False
     
@@ -301,7 +321,7 @@ class AuthManager:
                 
                 # Try to get user info (this endpoint may not exist in the actual API)
                 response = client.get(
-                    f"{self.config.api_url}/api/user",
+                    f"{self.config.api_url}/user",
                     headers=headers,
                     timeout=10.0
                 )
@@ -328,7 +348,7 @@ class AuthManager:
                 }
                 
                 response = client.get(
-                    f"{self.config.api_url}/api/orgs",
+                    f"{self.config.api_url}/orgs",
                     headers=headers,
                     timeout=10.0
                 )

alphai/cleanup.py

@@ -0,0 +1,351 @@
+"""Cleanup management for alphai CLI.
+
+Provides a class-based approach to resource cleanup that replaces
+global mutable state and handles signal registration safely.
+"""
+
+import signal
+import sys
+from typing import Optional, Callable, Dict, Any
+from rich.console import Console
+
+from .utils import get_logger
+
+logger = get_logger(__name__)
+
+
+class CleanupManager:
+    """Manages cleanup of resources with proper signal handling.
+    
+    This class replaces the global _cleanup_state and _jupyter_cleanup_state
+    dictionaries with a proper encapsulated approach that:
+    - Avoids global mutable state
+    - Allows safe signal handler installation/restoration
+    - Can be tested in isolation
+    """
+    
+    def __init__(self, console: Optional[Console] = None):
+        """Initialize the cleanup manager.
+        
+        Args:
+            console: Rich console for output. Creates one if not provided.
+        """
+        self.console = console or Console()
+        self._resources: Dict[str, Dict[str, Any]] = {}
+        self._cleanup_done = False
+        self._original_handlers: Dict[int, Any] = {}
+        self._cleanup_callbacks: list[Callable[[], None]] = []
+    
+    def register_resource(
+        self,
+        resource_type: str,
+        resource_id: str,
+        cleanup_fn: Optional[Callable[[], bool]] = None,
+        **metadata
+    ) -> None:
+        """Register a resource for cleanup.
+        
+        Args:
+            resource_type: Type of resource (e.g., 'container', 'tunnel', 'process')
+            resource_id: Unique identifier for the resource
+            cleanup_fn: Optional function to call for cleanup (returns success bool)
+            **metadata: Additional metadata to store with the resource
+        """
+        key = f"{resource_type}:{resource_id}"
+        self._resources[key] = {
+            'type': resource_type,
+            'id': resource_id,
+            'cleanup_fn': cleanup_fn,
+            **metadata
+        }
+        logger.debug(f"Registered resource for cleanup: {key}")
+    
+    def unregister_resource(self, resource_type: str, resource_id: str) -> None:
+        """Unregister a resource (e.g., after successful manual cleanup).
+        
+        Args:
+            resource_type: Type of resource
+            resource_id: Unique identifier for the resource
+        """
+        key = f"{resource_type}:{resource_id}"
+        if key in self._resources:
+            del self._resources[key]
+            logger.debug(f"Unregistered resource: {key}")
+    
+    def add_cleanup_callback(self, callback: Callable[[], None]) -> None:
+        """Add a cleanup callback to be called during cleanup.
+        
+        Args:
+            callback: Function to call during cleanup (no arguments, no return)
+        """
+        self._cleanup_callbacks.append(callback)
+    
+    def get_resource(self, resource_type: str, resource_id: str) -> Optional[Dict[str, Any]]:
+        """Get a registered resource by type and ID.
+        
+        Args:
+            resource_type: Type of resource
+            resource_id: Unique identifier for the resource
+            
+        Returns:
+            Resource metadata dict or None if not found
+        """
+        key = f"{resource_type}:{resource_id}"
+        return self._resources.get(key)
+    
+    def has_resources(self) -> bool:
+        """Check if there are any resources registered for cleanup."""
+        return len(self._resources) > 0 or len(self._cleanup_callbacks) > 0
+    
+    def install_signal_handlers(self) -> None:
+        """Install signal handlers for cleanup on interrupt.
+        
+        Stores original handlers so they can be restored later.
+        """
+        if self._original_handlers:
+            logger.debug("Signal handlers already installed")
+            return
+        
+        # Store original handlers
+        self._original_handlers[signal.SIGINT] = signal.signal(
+            signal.SIGINT, self._signal_handler
+        )
+        self._original_handlers[signal.SIGTERM] = signal.signal(
+            signal.SIGTERM, self._signal_handler
+        )
+        logger.debug("Signal handlers installed")
+    
+    def restore_signal_handlers(self) -> None:
+        """Restore original signal handlers."""
+        for sig, handler in self._original_handlers.items():
+            signal.signal(sig, handler)
+        self._original_handlers.clear()
+        logger.debug("Signal handlers restored")
+    
+    def _signal_handler(self, signum: int, frame) -> None:
+        """Handle interrupt signals by running cleanup.
+        
+        Args:
+            signum: Signal number received
+            frame: Current stack frame
+        """
+        if self._cleanup_done:
+            return
+        
+        if not self.has_resources():
+            sys.exit(0)
+        
+        self.cleanup()
+        sys.exit(0)
+    
+    def cleanup(self) -> bool:
+        """Run cleanup for all registered resources.
+        
+        Returns:
+            True if all cleanup succeeded, False if any failed
+        """
+        if self._cleanup_done:
+            logger.debug("Cleanup already done, skipping")
+            return True
+        
+        if not self.has_resources():
+            logger.debug("No resources to clean up")
+            return True
+        
+        self.console.print("\n[yellow]🔄 Cleaning up resources...[/yellow]")
+        success = True
+        
+        # Run cleanup callbacks first
+        for callback in self._cleanup_callbacks:
+            try:
+                callback()
+            except Exception as e:
+                logger.error(f"Cleanup callback failed: {e}", exc_info=True)
+                self.console.print(f"[red]Error in cleanup callback: {e}[/red]")
+                success = False
+        
+        # Clean up registered resources
+        for key, resource in list(self._resources.items()):
+            cleanup_fn = resource.get('cleanup_fn')
+            if cleanup_fn:
+                try:
+                    logger.debug(f"Running cleanup for {key}")
+                    if not cleanup_fn():
+                        success = False
+                except Exception as e:
+                    logger.error(f"Cleanup failed for {key}: {e}", exc_info=True)
+                    self.console.print(f"[red]Error cleaning up {resource['type']}: {e}[/red]")
+                    success = False
+        
+        if success:
+            self.console.print("[green]✓ Cleanup completed[/green]")
+        else:
+            self.console.print("[yellow]⚠ Cleanup completed with warnings[/yellow]")
+        
+        self._cleanup_done = True
+        self._resources.clear()
+        self._cleanup_callbacks.clear()
+        
+        return success
+    
+    def reset(self) -> None:
+        """Reset the cleanup manager state (useful for testing)."""
+        self._resources.clear()
+        self._cleanup_callbacks.clear()
+        self._cleanup_done = False
+        self.restore_signal_handlers()
+
+
+class DockerCleanupManager(CleanupManager):
+    """Cleanup manager specialized for Docker container cleanup."""
+    
+    def __init__(
+        self,
+        console: Optional[Console] = None,
+        docker_manager: Any = None,
+        client: Any = None
+    ):
+        """Initialize Docker cleanup manager.
+        
+        Args:
+            console: Rich console for output
+            docker_manager: DockerManager instance for container operations
+            client: AlphAIClient instance for API operations
+        """
+        super().__init__(console)
+        self.docker_manager = docker_manager
+        self.client = client
+        self._container_id: Optional[str] = None
+        self._tunnel_id: Optional[str] = None
+        self._project_id: Optional[str] = None
+    
+    def set_container(self, container_id: str) -> None:
+        """Set the container ID for cleanup."""
+        self._container_id = container_id
+    
+    def set_tunnel(self, tunnel_id: str) -> None:
+        """Set the tunnel ID for cleanup."""
+        self._tunnel_id = tunnel_id
+    
+    def set_project(self, project_id: str) -> None:
+        """Set the project ID for cleanup."""
+        self._project_id = project_id
+    
+    def cleanup(self) -> bool:
+        """Clean up Docker resources (container, tunnel, project)."""
+        if self._cleanup_done:
+            return True
+        
+        if not any([self._container_id, self._tunnel_id, self._project_id]):
+            return True
+        
+        self.console.print("\n[yellow]🔄 Cleaning up resources...[/yellow]")
+        success = True
+        
+        try:
+            # Clean up container and cloudflared service
+            if self._container_id and self.docker_manager:
+                if not self.docker_manager.cleanup_container_and_tunnel(
+                    container_id=self._container_id,
+                    tunnel_id=self._tunnel_id,
+                    project_id=self._project_id,
+                    force=True
+                ):
+                    success = False
+            
+            # Clean up tunnel and project via API
+            if self.client and (self._tunnel_id or self._project_id):
+                if not self.client.cleanup_tunnel_and_project(
+                    tunnel_id=self._tunnel_id,
+                    project_id=self._project_id,
+                    force=True
+                ):
+                    success = False
+            
+            if success:
+                self.console.print("[green]✓ Cleanup completed[/green]")
+            else:
+                self.console.print("[yellow]⚠ Cleanup completed with warnings[/yellow]")
+                
+        except Exception as e:
+            logger.error(f"Error during cleanup: {e}", exc_info=True)
+            self.console.print(f"[red]Error during cleanup: {e}[/red]")
+            success = False
+        
+        # Reset state
+        self._container_id = None
+        self._tunnel_id = None
+        self._project_id = None
+        self._cleanup_done = True
+        
+        return success
+
+
+class JupyterCleanupManager(CleanupManager):
+    """Cleanup manager specialized for Jupyter process cleanup."""
+    
+    def __init__(
+        self,
+        console: Optional[Console] = None,
+        jupyter_manager: Any = None,
+        client: Any = None
+    ):
+        """Initialize Jupyter cleanup manager.
+        
+        Args:
+            console: Rich console for output
+            jupyter_manager: JupyterManager instance for process operations
+            client: AlphAIClient instance for API operations
+        """
+        super().__init__(console)
+        self.jupyter_manager = jupyter_manager
+        self.client = client
+        self._tunnel_id: Optional[str] = None
+        self._project_id: Optional[str] = None
+    
+    def set_tunnel(self, tunnel_id: str) -> None:
+        """Set the tunnel ID for cleanup."""
+        self._tunnel_id = tunnel_id
+    
+    def set_project(self, project_id: str) -> None:
+        """Set the project ID for cleanup."""
+        self._project_id = project_id
+    
+    def cleanup(self) -> bool:
+        """Clean up Jupyter resources (process, tunnel, project)."""
+        if self._cleanup_done:
+            return True
+        
+        if not self.jupyter_manager:
+            return True
+        
+        self.console.print("\n[yellow]🔄 Cleaning up Jupyter resources...[/yellow]")
+        success = True
+        
+        try:
+            # Comprehensive cleanup via JupyterManager
+            if not self.jupyter_manager.cleanup(
+                client=self.client,
+                tunnel_id=self._tunnel_id,
+                project_id=self._project_id,
+                force=True
+            ):
+                success = False
+            
+            if success:
+                self.console.print("[green]✓ Cleanup completed[/green]")
+            else:
+                self.console.print("[yellow]⚠ Cleanup completed with warnings[/yellow]")
+                
+        except Exception as e:
+            logger.error(f"Error during Jupyter cleanup: {e}", exc_info=True)
+            self.console.print(f"[red]Error during cleanup: {e}[/red]")
+            success = False
+        
+        # Reset state
+        self._tunnel_id = None
+        self._project_id = None
+        self._cleanup_done = True
+        
+        return success
+