gitarsenal-cli/lib/e2b-sandbox.js

CLI tool for creating Modal sandboxes with GitHub repositories

const path = require('path');
const fs = require('fs-extra');
const { spawn } = require('child_process');
const chalk = require('chalk');
const ora = require('ora');
const { promisify } = require('util');
const { exec } = require('child_process');
const os = require('os');

const execAsync = promisify(exec);

/**
 * Get the Python executable path from the virtual environment
 * @returns {string} - Path to Python executable or 'python3'
 */
function getPythonExecutable() {
  // Check if PYTHON_EXECUTABLE is set by the virtual environment activation
  if (process.env.PYTHON_EXECUTABLE && fs.existsSync(process.env.PYTHON_EXECUTABLE)) {
    return process.env.PYTHON_EXECUTABLE;
  }
  
  // Try to find the virtual environment Python
  const venvPath = path.join(__dirname, '..', '.venv');
  const isWindows = process.platform === 'win32';
  
  // Check for uv-style virtual environment first
  const uvPythonPath = path.join(venvPath, 'bin', 'python');
  
  // Check for traditional venv structure
  const traditionalPythonPath = isWindows ? 
    path.join(venvPath, 'Scripts', 'python.exe') : 
    path.join(venvPath, 'bin', 'python');
  
  if (fs.existsSync(uvPythonPath)) {
    return uvPythonPath;
  } else if (fs.existsSync(traditionalPythonPath)) {
    return traditionalPythonPath;
  }
  
  // Fall back to system Python
  return isWindows ? 'python' : 'python3';
}

/**
 * Fetch Anthropic API key directly from the server
 * @returns {Promise<string>} - The Anthropic API key
 */
async function fetchAnthropicApiKey() {
  try {
    const pythonExecutable = getPythonExecutable();
    const fetchTokensScript = path.join(__dirname, '..', 'python', 'fetch_modal_tokens.py');
    
    // Run the fetch_modal_tokens.py script with --output-anthropic-key flag
    const { stdout, stderr } = await execAsync(`${pythonExecutable} ${fetchTokensScript} --output-anthropic-key`, {
      env: { ...process.env },
      timeout: 30000 // 30 seconds timeout
    });
    
    // Parse the output to find the Anthropic API key
    const anthropicKeyMatch = stdout.match(/ANTHROPIC_API_KEY=([a-zA-Z0-9\-_]+)/);
    if (anthropicKeyMatch && anthropicKeyMatch[1]) {
      const anthropicApiKey = anthropicKeyMatch[1];
      return anthropicApiKey;
    }
    
    return null;
  } catch (error) {
    return null;
  }
}

/**
 * Fetch E2B API key directly from the server
 * @returns {Promise<string>} - The E2B API key
 */
async function fetchE2BApiKey() {
  try {
    
    const pythonExecutable = getPythonExecutable();
    const fetchTokensScript = path.join(__dirname, '..', 'python', 'fetch_modal_tokens.py');
    
    // Run the fetch_modal_tokens.py script with --output-e2b-key flag
    const { stdout, stderr } = await execAsync(`${pythonExecutable} ${fetchTokensScript} --output-e2b-key`, {
      env: { ...process.env },
      timeout: 30000 // 30 seconds timeout
    });
    
    // Parse the output to find the E2B API key
    const e2bKeyMatch = stdout.match(/E2B_API_KEY=(e2b_[a-zA-Z0-9]+)/);
    if (e2bKeyMatch && e2bKeyMatch[1]) {
      const e2bApiKey = e2bKeyMatch[1];
      return e2bApiKey;
    }
    
    console.error(chalk.yellow('⚠️ Could not extract E2B API key from fetch_modal_tokens.py output'));
    console.error(chalk.gray('Output: ' + stdout.substring(0, 100) + '...'));
    return null;
  } catch (error) {
    console.error(chalk.red(`❌ Error fetching E2B API key: ${error.message}`));
    return null;
  }
}

/**
 * Run an E2B sandbox with the given options
 * @param {Object} options - Sandbox options
 * @param {string} options.repoUrl - GitHub repository URL
 * @param {Array<string>} options.setupCommands - Setup commands
 * @returns {Promise<void>}
 */
async function runE2BSandbox(options) {
  const { 
    repoUrl, 
    setupCommands = [], 
    userId,
    userName,
    userEmail,
    apiKeys = {}
  } = options;
  
  console.log(chalk.blue('🚀 Starting E2B sandbox...'));
  
  // Check if E2B_API_KEY is already set in environment variables
  let e2bApiKey = process.env.E2B_API_KEY;
  
  // Check if the key is a placeholder value
  if (e2bApiKey && (e2bApiKey === 'your_e2b_api_key' || e2bApiKey.startsWith('placeholder'))) {
    console.log(chalk.yellow('⚠️ Found placeholder E2B API key in environment, will try to get real key'));
    e2bApiKey = null;
  }
  
  // If not, try to get it from apiKeys
  if (!e2bApiKey && apiKeys.E2B_API_KEY) {
    e2bApiKey = apiKeys.E2B_API_KEY;
    // Check if the key from apiKeys is a placeholder
    if (e2bApiKey === 'your_e2b_api_key' || e2bApiKey.startsWith('placeholder')) {
      console.log(chalk.yellow('⚠️ Found placeholder E2B API key in apiKeys, will try to get real key'));
      e2bApiKey = null;
    } else {
      console.log(chalk.green('✅ Using E2B API key from apiKeys'));
    }
  }
  
  // If still not found, try to fetch it directly from the server
  if (!e2bApiKey) {
    e2bApiKey = await fetchE2BApiKey();
  }
  
  // If still not found, try to get it from modal_tokens.json
  if (!e2bApiKey) {
    try {
      const modalTokensPath = path.join(__dirname, '..', 'python', 'modal_tokens.json');
      if (fs.existsSync(modalTokensPath)) {
        const modalTokens = JSON.parse(fs.readFileSync(modalTokensPath, 'utf8'));
        if (modalTokens.e2b_api_key) {
          e2bApiKey = modalTokens.e2b_api_key;
          console.log(chalk.green('✅ Using E2B API key from modal_tokens.json'));
        }
      }
    } catch (error) {
      console.error(chalk.yellow(`⚠️ Error reading modal_tokens.json: ${error.message}`));
    }
  }

  // Check if we have a valid E2B API key
  if (!e2bApiKey) {
    console.error(chalk.red('❌ E2B API key not found'));
    console.error(chalk.yellow('Please set the E2B_API_KEY environment variable or add it to your API keys'));
    return { success: false, error: 'E2B API key not found' };
  }

  // Set up the environment for the Python process
  const pythonExecutable = getPythonExecutable();
  if (!pythonExecutable) {
    console.error(chalk.red('❌ Python executable not found'));
    return { success: false, error: 'Python executable not found' };
  }

  // Get OpenAI and Anthropic API keys from environment or apiKeys
  const openaiApiKey = apiKeys.OPENAI_API_KEY || process.env.OPENAI_API_KEY;
  let anthropicApiKey = apiKeys.ANTHROPIC_API_KEY || process.env.ANTHROPIC_API_KEY;
  
  // If Anthropic API key is not found, try to fetch it directly from the server
  if (!anthropicApiKey) {
    anthropicApiKey = await fetchAnthropicApiKey();
    if (anthropicApiKey) {
      // Add it to apiKeys so it gets passed to the Python script
      apiKeys.ANTHROPIC_API_KEY = anthropicApiKey;
    }
  }
  
  // Run the e2b_sandbox_agent.py script with the repository URL
  const scriptPath = path.join(__dirname, '..', 'python', 'e2b_sandbox_agent.py');
  
  // Prepare command line arguments
  const args = [
    scriptPath,
    '--repo', repoUrl
  ];
  
  // Add API keys if available
  if (openaiApiKey) {
    args.push('--openai-api-key', openaiApiKey);
  }
  
  if (anthropicApiKey) {
    args.push('--anthropic-api-key', anthropicApiKey);
  }

  // Run the Python script
  const pythonProcess = spawn(pythonExecutable, args, {
    stdio: 'inherit',
    env: {
      ...process.env,
      E2B_API_KEY: e2bApiKey
    }
  });

  // Clean up function to remove the E2B API key
  const cleanupE2BApiKey = () => {
    try {
      // Remove the E2B API key from modal_tokens.json if it exists
      const modalTokensPath = path.join(__dirname, '..', 'python', 'modal_tokens.json');
      if (fs.existsSync(modalTokensPath)) {
        const modalTokens = JSON.parse(fs.readFileSync(modalTokensPath, 'utf8'));
        if (modalTokens.e2b_api_key) {
          delete modalTokens.e2b_api_key;
          fs.writeFileSync(modalTokensPath, JSON.stringify(modalTokens, null, 2));
          console.log(chalk.green('🧹 Removed E2B API key from modal_tokens.json'));
        }
      }
    } catch (error) {
      console.error(chalk.yellow(`⚠️ Error cleaning up E2B API key: ${error.message}`));
    }
  };

  return new Promise((resolve, reject) => {
    pythonProcess.on('close', (code) => {
      // Clean up the E2B API key
      cleanupE2BApiKey();
      
      if (code === 0 || code === 130) {  // 130 is the exit code for SIGINT (Ctrl+C)
        console.log(chalk.green('✅ E2B sandbox session completed'));
        resolve({ success: true });
      } else {
        console.error(chalk.red(`✗ E2B sandbox session failed`));
        console.error(chalk.red(`  Exit code: ${code}`));
        reject(new Error(`Process exited with code ${code}`));
      }
    });

    pythonProcess.on('error', (error) => {
      // Clean up the E2B API key
      cleanupE2BApiKey();
      
      console.error(chalk.red(`✗ E2B sandbox session failed`));
      console.error(chalk.red(`  Error: ${error.message}`));
      reject(error);
    });
    
    // Handle process termination signals
    const handleTermination = () => {
      if (pythonProcess && !pythonProcess.killed) {
        pythonProcess.kill('SIGINT');
      }
      cleanupE2BApiKey();
    };
    
    // Listen for termination signals
    process.on('SIGINT', handleTermination);
    process.on('SIGTERM', handleTermination);
    process.on('exit', handleTermination);
  });
}

module.exports = {
  runE2BSandbox
};