suna/backend/utils/scripts/archive_old_sandboxes.py

#!/usr/bin/env python
"""
Script to archive sandboxes for projects that are older than 1 day.

Usage:
    python archive_old_sandboxes.py [--days N] [--dry-run]

This script:
1. Gets all projects from the projects table
2. Filters projects created more than N days ago (default: 1 day)
3. Archives the sandboxes for those projects

Make sure your environment variables are properly set:
- SUPABASE_URL
- SUPABASE_SERVICE_ROLE_KEY
- DAYTONA_SERVER_URL
"""

# TODO: SAVE THE LATEST SANDBOX STATE SOMEWHERE OR LIKE MASS CHECK THE STATE BEFORE STARTING TO ARCHIVE - AS ITS GOING TO GO OVER A BUNCH THAT ARE ALREADY ARCHIVED – MAYBE BEST TO GET ALL FROM DAYTONA AND THEN RUN THE ARCHIVE ONLY ON THE ONES THAT MEET THE CRITERIA (STOPPED STATE)

import asyncio
import sys
import os
import argparse
from typing import List, Dict, Any
from datetime import datetime, timedelta
from dotenv import load_dotenv

# Load script-specific environment variables
load_dotenv(".env")

from services.supabase import DBConnection
from sandbox.sandbox import daytona
from utils.logger import logger

# Global DB connection to reuse
db_connection = None


async def get_old_projects(days_threshold: int = 1) -> List[Dict[str, Any]]:
    """
    Query all projects created more than N days ago.
    
    Args:
        days_threshold: Number of days threshold (default: 1)
        
    Returns:
        List of projects with their sandbox information
    """
    global db_connection
    if db_connection is None:
        db_connection = DBConnection()
    
    client = await db_connection.client
    
    # Print the Supabase URL being used
    print(f"Using Supabase URL: {os.getenv('SUPABASE_URL')}")
    
    # Calculate the date threshold
    threshold_date = (datetime.now() - timedelta(days=days_threshold)).isoformat()
    
    # Initialize variables for pagination
    all_projects = []
    page_size = 1000
    current_page = 0
    has_more = True
    
    logger.info(f"Starting to fetch projects older than {days_threshold} day(s)")
    print(f"Looking for projects created before: {threshold_date}")
    
    # Paginate through all projects
    while has_more:
        # Query projects with pagination
        start_range = current_page * page_size
        end_range = start_range + page_size - 1
        
        logger.info(f"Fetching projects page {current_page+1} (range: {start_range}-{end_range})")
        
        try:
            result = await client.table('projects').select(
                'project_id',
                'name',
                'created_at',
                'account_id',
                'sandbox'
            ).order('created_at', desc=True).range(start_range, end_range).execute()
            
            # Debug info - print raw response
            print(f"Response data length: {len(result.data)}")
            
            if not result.data:
                print("No more data returned from query, ending pagination")
                has_more = False
            else:
                # Print a sample project to see the actual data structure
                if current_page == 0 and result.data:
                    print(f"Sample project data: {result.data[0]}")
                
                all_projects.extend(result.data)
                current_page += 1
                
                # Progress update
                logger.info(f"Loaded {len(all_projects)} projects so far")
                print(f"Loaded {len(all_projects)} projects so far...")
                
                # Check if we've reached the end - if we got fewer results than the page size
                if len(result.data) < page_size:
                    print(f"Got {len(result.data)} records which is less than page size {page_size}, ending pagination")
                    has_more = False
                else:
                    print(f"Full page returned ({len(result.data)} records), continuing to next page")
                    
        except Exception as e:
            logger.error(f"Error during pagination: {str(e)}")
            print(f"Error during pagination: {str(e)}")
            has_more = False  # Stop on error
    
    # Print the query result summary
    total_projects = len(all_projects)
    print(f"Found {total_projects} total projects in database")
    logger.info(f"Total projects found in database: {total_projects}")
    
    if not all_projects:
        logger.info("No projects found in database")
        return []
    
    # Filter projects that are older than the threshold and have sandbox information
    old_projects_with_sandboxes = [
        project for project in all_projects
        if project.get('created_at') and project.get('created_at') < threshold_date
        and project.get('sandbox') and project['sandbox'].get('id')
    ]
    
    logger.info(f"Found {len(old_projects_with_sandboxes)} old projects with sandboxes")
    
    # Print a few sample old projects for debugging
    if old_projects_with_sandboxes:
        print("\nSample of old projects with sandboxes:")
        for i, project in enumerate(old_projects_with_sandboxes[:3]):
            print(f"  {i+1}. {project.get('name')} (Created: {project.get('created_at')})")
            print(f"     Sandbox ID: {project['sandbox'].get('id')}")
            if i >= 2:
                break
    
    return old_projects_with_sandboxes


async def archive_sandbox(project: Dict[str, Any], dry_run: bool) -> bool:
    """
    Archive a single sandbox.
    
    Args:
        project: Project information containing sandbox to archive
        dry_run: If True, only simulate archiving
        
    Returns:
        True if successful, False otherwise
    """
    sandbox_id = project['sandbox'].get('id')
    project_name = project.get('name', 'Unknown')
    project_id = project.get('project_id', 'Unknown')
    created_at = project.get('created_at', 'Unknown')
    
    try:
        logger.info(f"Checking sandbox {sandbox_id} for project '{project_name}' (ID: {project_id}, Created: {created_at})")
        
        if dry_run:
            logger.info(f"DRY RUN: Would archive sandbox {sandbox_id}")
            print(f"Would archive sandbox {sandbox_id} for project '{project_name}' (Created: {created_at})")
            return True
        
        # Get the sandbox
        sandbox = daytona.get(sandbox_id)
        
        # Check sandbox state - it must be stopped before archiving
        sandbox_info = sandbox.info()
        
        # Log the current state
        logger.info(f"Sandbox {sandbox_id} is in '{sandbox_info.state}' state")
        
        # Only archive if the sandbox is in the stopped state
        if sandbox_info.state == "stopped":
            logger.info(f"Archiving sandbox {sandbox_id} as it is in stopped state")
            sandbox.archive()
            logger.info(f"Successfully archived sandbox {sandbox_id}")
            return True
        else:
            logger.info(f"Skipping sandbox {sandbox_id} as it is not in stopped state (current: {sandbox_info.state})")
            return True
            
    except Exception as e:
        import traceback
        error_type = type(e).__name__
        stack_trace = traceback.format_exc()
        
        # Log detailed error information
        logger.error(f"Error processing sandbox {sandbox_id}: {str(e)}")
        logger.error(f"Error type: {error_type}")
        logger.error(f"Stack trace:\n{stack_trace}")
        
        # If the exception has a response attribute (like in HTTP errors), log it
        if hasattr(e, 'response'):
            try:
                response_data = e.response.json() if hasattr(e.response, 'json') else str(e.response)
                logger.error(f"Response data: {response_data}")
            except Exception:
                logger.error(f"Could not parse response data from error")
        
        print(f"Failed to process sandbox {sandbox_id}: {error_type} - {str(e)}")
        return False


async def process_sandboxes(old_projects: List[Dict[str, Any]], dry_run: bool) -> tuple[int, int]:
    """
    Process all sandboxes sequentially.
    
    Args:
        old_projects: List of projects older than the threshold
        dry_run: Whether to actually archive sandboxes or just simulate
        
    Returns:
        Tuple of (processed_count, failed_count)
    """
    processed_count = 0
    failed_count = 0
    
    if dry_run:
        logger.info(f"DRY RUN: Would archive {len(old_projects)} sandboxes")
    else:
        logger.info(f"Archiving {len(old_projects)} sandboxes")
    
    print(f"Processing {len(old_projects)} sandboxes...")
    
    # Process each sandbox sequentially
    for i, project in enumerate(old_projects):
        success = await archive_sandbox(project, dry_run)
        
        if success:
            processed_count += 1
        else:
            failed_count += 1
        
        # Print progress periodically
        if (i + 1) % 20 == 0 or (i + 1) == len(old_projects):
            progress = (i + 1) / len(old_projects) * 100
            print(f"Progress: {i + 1}/{len(old_projects)} sandboxes processed ({progress:.1f}%)")
            print(f"  - Processed: {processed_count}, Failed: {failed_count}")
    
    return processed_count, failed_count


async def main():
    """Main function to run the script."""
    # Parse command line arguments
    parser = argparse.ArgumentParser(description='Archive sandboxes for projects older than N days')
    parser.add_argument('--days', type=int, default=1, help='Age threshold in days (default: 1)')
    parser.add_argument('--dry-run', action='store_true', help='Show what would be archived without actually archiving')
    args = parser.parse_args()

    logger.info(f"Starting sandbox cleanup for projects older than {args.days} day(s)")
    if args.dry_run:
        logger.info("DRY RUN MODE - No sandboxes will be archived")
    
    # Print environment info
    print(f"Environment Mode: {os.getenv('ENV_MODE', 'Not set')}")
    print(f"Daytona Server: {os.getenv('DAYTONA_SERVER_URL', 'Not set')}")
    
    try:
        # Initialize global DB connection
        global db_connection
        db_connection = DBConnection()
        
        # Get all projects older than the threshold
        old_projects = await get_old_projects(args.days)
        
        if not old_projects:
            logger.info(f"No projects older than {args.days} day(s) with sandboxes to process")
            print(f"No projects older than {args.days} day(s) with sandboxes to archive.")
            return
        
        # Print summary of what will be processed
        print("\n===== SANDBOX CLEANUP SUMMARY =====")
        print(f"Projects older than {args.days} day(s): {len(old_projects)}")
        print(f"Sandboxes that will be archived: {len(old_projects)}")
        print("===================================")
        
        logger.info(f"Found {len(old_projects)} projects older than {args.days} day(s)")
        
        # Ask for confirmation before proceeding
        if not args.dry_run:
            print("\n⚠️  WARNING: You are about to archive sandboxes for old projects ⚠️")
            print("This action cannot be undone!")
            confirmation = input("\nAre you sure you want to proceed with archiving? (TRUE/FALSE): ").strip().upper()
            
            if confirmation != "TRUE":
                print("Archiving cancelled. Exiting script.")
                logger.info("Archiving cancelled by user")
                return
            
            print("\nProceeding with sandbox archiving...\n")
            logger.info("User confirmed sandbox archiving")
        
        # List a sample of projects to be processed
        for i, project in enumerate(old_projects[:5]):  # Just show first 5 for brevity
            created_at = project.get('created_at', 'Unknown')
            project_name = project.get('name', 'Unknown')
            project_id = project.get('project_id', 'Unknown')
            sandbox_id = project['sandbox'].get('id')
            
            print(f"{i+1}. Project: {project_name}")
            print(f"   Project ID: {project_id}")
            print(f"   Created At: {created_at}")
            print(f"   Sandbox ID: {sandbox_id}")
            
        if len(old_projects) > 5:
            print(f"   ... and {len(old_projects) - 5} more projects")
        
        # Process all sandboxes
        processed_count, failed_count = await process_sandboxes(old_projects, args.dry_run)
        
        # Print final summary
        print("\nSandbox Cleanup Summary:")
        print(f"Total projects older than {args.days} day(s): {len(old_projects)}")
        print(f"Total sandboxes processed: {len(old_projects)}")
        
        if args.dry_run:
            print(f"DRY RUN: No sandboxes were actually archived")
        else:
            print(f"Successfully processed: {processed_count}")
            print(f"Failed to process: {failed_count}")
        
        logger.info("Sandbox cleanup completed")
            
    except Exception as e:
        logger.error(f"Error during sandbox cleanup: {str(e)}")
        sys.exit(1)
    finally:
        # Clean up database connection
        if db_connection:
            await DBConnection.disconnect()


if __name__ == "__main__":
    asyncio.run(main())