gitarsenal-cli 1.9.75 → 1.9.77

package/python/test_modalSandboxScript.py

@@ -11,48 +11,6 @@ import string
 import argparse
 from pathlib import Path
 import modal
-from auth_manager import AuthManager
-
-# Removed unused boxed output functions since they're no longer used with the Agent-based approach
-
-
-# Removed _execute_with_box function as it's no longer used with the Agent-based approach
-
-# Early argument parsing for proxy settings only
-early_parser = argparse.ArgumentParser(add_help=False)
-early_parser.add_argument('--proxy-url', help='URL of the proxy server')
-early_parser.add_argument('--proxy-api-key', help='API key for the proxy server')
-
-# Parse only proxy args early to avoid conflicts 
-early_args, _ = early_parser.parse_known_args()
-
-# Set proxy URL and API key in environment variables if provided
-if early_args.proxy_url:
-    os.environ["MODAL_PROXY_URL"] = early_args.proxy_url
-    
-if early_args.proxy_api_key:
-    os.environ["MODAL_PROXY_API_KEY"] = early_args.proxy_api_key
-
-# Import the fetch_modal_tokens module
-from fetch_modal_tokens import get_tokens
-token_id, token_secret, openai_api_key, anthropic_api_key, openrouter_api_key, groq_api_key = get_tokens()
-
-# Check if we got valid tokens
-if token_id is None or token_secret is None:
-    raise ValueError("Could not get valid tokens")
-
-# Explicitly set the environment variables again to be sure
-os.environ["MODAL_TOKEN_ID"] = token_id
-os.environ["MODAL_TOKEN_SECRET"] = token_secret
-if openai_api_key:
-    os.environ["OPENAI_API_KEY"] = openai_api_key
-if anthropic_api_key:
-    os.environ["ANTHROPIC_API_KEY"] = anthropic_api_key
-# Also set the old environment variable for backward compatibility
-os.environ["MODAL_TOKEN"] = token_id
-
-# Set token variables for later use
-token = token_id  # For backward compatibility
 
 
 def generate_random_password(length=16):
@@ -80,9 +38,281 @@ def get_stored_credentials():
         return {}
 
 
+# Global SSH container function (must be at global scope for Modal)
+def ssh_container_function(ssh_password=None, repo_url=None, repo_name=None, setup_commands=None, openai_api_key=None, anthropic_api_key=None, stored_credentials=None):
+    """Start SSH container with password authentication and intelligent repository setup using Agent."""
+    import subprocess
+    import time
+    import os
+    import json
+    import sys
+    
+    # Set root password
+    subprocess.run(["bash", "-c", f"echo 'root:{ssh_password}' | chpasswd"], check=True)
+    
+    # Set OpenAI API key if provided
+    if openai_api_key:
+        os.environ['OPENAI_API_KEY'] = openai_api_key
+    else:
+        print("⚠️ No OpenAI API key provided to container")
+    
+    # Set up stored credentials in container environment
+    if stored_credentials:
+        print(f"🔐 Setting up {len(stored_credentials)} stored credentials in container...")
+        for key, value in stored_credentials.items():
+            # Set each credential as an environment variable
+            env_var_name = key.upper().replace('-', '_').replace(' ', '_')
+            os.environ[env_var_name] = value
+            print(f"✅ Set {env_var_name} in container environment")
+        
+        # Also save credentials to a file in the container for easy access
+        credentials_dir = "/root/.gitarsenal"
+        os.makedirs(credentials_dir, exist_ok=True)
+        credentials_file = os.path.join(credentials_dir, "credentials.json")
+        with open(credentials_file, 'w') as f:
+            json.dump(stored_credentials, f, indent=2)
+        print(f"✅ Saved credentials to {credentials_file}")
+        
+        # Print available credentials for user reference
+        print("\n🔐 AVAILABLE CREDENTIALS IN CONTAINER:")
+        print("="*50)
+        for key, value in stored_credentials.items():
+            masked_value = value[:8] + "..." if len(value) > 8 else "***"
+            env_var_name = key.upper().replace('-', '_').replace(' ', '_')
+            print(f"  {key} -> {env_var_name} = {masked_value}")
+        print("="*50)
+        print("💡 These credentials are available as environment variables and in /root/.gitarsenal/credentials.json")
+        
+    
+    # Start SSH service
+    subprocess.run(["service", "ssh", "start"], check=True)
+    
+    # Use Agent for intelligent repository setup
+    if repo_url:
+        print("🤖 Using Agent for intelligent repository setup...")
+        
+        # Set up environment variables for the Agent
+        if openai_api_key:
+            os.environ['OPENAI_API_KEY'] = openai_api_key
+        if anthropic_api_key:
+            os.environ['ANTHROPIC_API_KEY'] = anthropic_api_key
+        
+        # Set up Anthropic API key from stored credentials
+        if stored_credentials:
+            # Look for Anthropic API key in various possible names
+            for key_name in ['ANTHROPIC_API_KEY', 'anthropic_api_key', 'anthropic-api-key']:
+                if key_name in stored_credentials:
+                    anthropic_api_key = stored_credentials[key_name]
+                    os.environ['ANTHROPIC_API_KEY'] = anthropic_api_key
+                    print(f"✅ Set Anthropic API key from stored credentials")
+                    break
+                    
+            if not anthropic_api_key:
+                print("⚠️ No Anthropic API key found in stored credentials")
+                print("💡 Agent will require an Anthropic API key for operation")
+        
+        try:
+            print("🔧 Running Agent for repository setup...")
+            
+            print("\n" + "="*80)
+            print("🤖 AGENT REPOSITORY SETUP")
+            print("="*80)
+            print(f"Repository: {repo_url}")
+            print(f"Working Directory: /root")
+            if stored_credentials:
+                print(f"Available Credentials: {len(stored_credentials)} items")
+            print("="*80 + "\n")
+            
+            # Call Agent directly as subprocess with real-time output
+            claude_prompt = f"clone, setup and run {repo_url}"
+            print(f"🚀 Executing the task: \"{claude_prompt}\"")
+            print("\n" + "="*60)
+            print("🎉 AGENT OUTPUT (LIVE)")
+            print("="*60)
+            
+            # Use Popen for real-time output streaming with optimizations
+            import sys
+            import select
+            import fcntl
+            import os as os_module
+            
+            process = subprocess.Popen(
+                ["python", "-u", "/python/kill_claude/claude_code_agent.py", claude_prompt],  # -u for unbuffered output
+                cwd="/root",
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,  # Keep separate for better handling
+                text=True,
+                bufsize=0,  # Unbuffered for fastest output
+                universal_newlines=True,
+                env=dict(os.environ, PYTHONUNBUFFERED='1')  # Force unbuffered Python output
+            )
+            
+            # Make stdout and stderr non-blocking for faster reading
+            def make_non_blocking(fd):
+                flags = fcntl.fcntl(fd, fcntl.F_GETFL)
+                fcntl.fcntl(fd, fcntl.F_SETFL, flags | os_module.O_NONBLOCK)
+            
+            make_non_blocking(process.stdout)
+            make_non_blocking(process.stderr)
+            
+            # Stream output in real-time with robust error handling
+            try:
+                stdout_buffer = ""
+                stderr_buffer = ""
+                
+                while process.poll() is None:
+                    try:
+                        # Use select for efficient I/O multiplexing with error handling
+                        ready, _, _ = select.select([process.stdout, process.stderr], [], [], 0.1)  # 100ms timeout
+                        
+                        for stream in ready:
+                            try:
+                                if stream == process.stdout:
+                                    chunk = stream.read(1024)  # Read in chunks for efficiency
+                                    if chunk is not None and chunk:
+                                        stdout_buffer += chunk
+                                        # Process complete lines immediately
+                                        while '\n' in stdout_buffer:
+                                            line, stdout_buffer = stdout_buffer.split('\n', 1)
+                                            print(line, flush=True)  # Force immediate flush
+                                elif stream == process.stderr:
+                                    chunk = stream.read(1024)
+                                    if chunk is not None and chunk:
+                                        stderr_buffer += chunk
+                                        # Process complete lines immediately  
+                                        while '\n' in stderr_buffer:
+                                            line, stderr_buffer = stderr_buffer.split('\n', 1)
+                                            print(f"STDERR: {line}", flush=True)
+                            except (BlockingIOError, OSError, ValueError):
+                                # Handle various I/O errors gracefully
+                                continue
+                    except (select.error, OSError):
+                        # If select fails, fall back to simple polling
+                        time.sleep(0.1)
+                        continue
+                
+                # Process any remaining output after process ends
+                try:
+                    # Read any remaining data from streams
+                    remaining_stdout = process.stdout.read()
+                    remaining_stderr = process.stderr.read()
+                    
+                    if remaining_stdout:
+                        stdout_buffer += remaining_stdout
+                    if remaining_stderr:
+                        stderr_buffer += remaining_stderr
+                    
+                    # Output remaining buffered content
+                    if stdout_buffer.strip():
+                        print(stdout_buffer.strip(), flush=True)
+                    if stderr_buffer.strip():
+                        print(f"STDERR: {stderr_buffer.strip()}", flush=True)
+                except (OSError, ValueError):
+                    # Handle cases where streams are already closed
+                    pass
+                
+                # Get final return code
+                return_code = process.returncode
+                
+                print("\n" + "="*60)
+                if return_code == 0:
+                    print("✅ Agent completed successfully!")
+                else:
+                    print(f"⚠️ Agent exited with code: {return_code}")
+                print("="*60)
+                
+            except subprocess.TimeoutExpired:
+                print("\n⚠️ Agent timed out after 10 minutes")
+                process.kill()
+                process.wait()
+            except Exception as stream_error:
+                pass
+                
+                # Fallback to simple readline approach
+                try:
+                    # Restart the process with simpler streaming
+                    if process.poll() is None:
+                        process.kill()
+                        process.wait()
+                    
+                    fallback_process = subprocess.Popen(
+                        ["python", "-u", "/python/kill_claude/claude_code_agent.py", claude_prompt],
+                        cwd="/root",
+                        stdout=subprocess.PIPE,
+                        stderr=subprocess.STDOUT,
+                        text=True,
+                        bufsize=1,
+                        universal_newlines=True
+                    )
+                    
+                    # Simple line-by-line reading
+                    while True:
+                        line = fallback_process.stdout.readline()
+                        if line == '' and fallback_process.poll() is not None:
+                            break
+                        if line:
+                            print(line.rstrip(), flush=True)
+                    
+                    return_code = fallback_process.returncode
+                    
+                    print("\n" + "="*60)
+                    if return_code == 0:
+                        print("✅ Agent completed successfully!")
+                    else:
+                        print(f"⚠️ Agent exited with code: {return_code}")
+                    print("="*60)
+                    
+                except Exception as fallback_error:
+                    print(f"\n❌ Fallback streaming also failed: {fallback_error}")
+                    print("⚠️ Agent may have completed, but output streaming failed")
+                    return_code = 1
+            
+        except Exception as e:
+            print(f"❌ Error during repository setup: {e}")
+            print("⚠️ Proceeding without setup...")
+            import traceback
+            traceback.print_exc()
+    else:
+        print("⚠️ No repository URL provided, skipping setup")
+    
+    print("🔌 Creating SSH tunnel on port 22...")
+    # Create SSH tunnel
+    with modal.forward(22, unencrypted=True) as tunnel:
+        host, port = tunnel.tcp_socket
+        
+        print("\n" + "=" * 80)
+        print("🎉 SSH CONTAINER IS READY!")
+        print("=" * 80)
+        print(f"🌐 SSH Host: {host}")
+        print(f"🔌 SSH Port: {port}")
+        print(f"👤 Username: root")
+        print(f"🔐 Password: {ssh_password}")
+        print()
+        print("🔗 CONNECT USING THIS COMMAND:")
+        print(f"ssh -p {port} root@{host}")
+        print("=" * 80)
+        
+        print("🔄 Starting keep-alive loop...")
+        # Keep the container running
+        iteration = 0
+        while True:
+            iteration += 1
+            if iteration % 10 == 1:  # Print every 5 minutes (10 * 30 seconds = 5 minutes)
+                print(f"💓 Container alive (iteration {iteration})")
+            
+            time.sleep(30)
+            # Check if SSH service is still running
+            try:
+                subprocess.run(["service", "ssh", "status"], check=True, 
+                             capture_output=True)
+            except subprocess.CalledProcessError:
+                print("⚠️ SSH service stopped, restarting...")
+                subprocess.run(["service", "ssh", "start"], check=True)
+
+
 # Create Modal SSH container with GPU support and intelligent repository setup using Agent
 def create_modal_ssh_container(gpu_type, repo_url=None, repo_name=None, setup_commands=None, 
-                               volume_name=None, timeout_minutes=60, ssh_password=None, interactive=False, gpu_count=1, use_cuda_base=False):
+                               volume_name=None, timeout_minutes=60, ssh_password=None, interactive=False, gpu_count=1):
     """Create a Modal SSH container with GPU support and intelligent repository setup.
     
     When repo_url is provided, uses Agent for intelligent repository setup.
@@ -240,25 +470,20 @@ def create_modal_ssh_container(gpu_type, repo_url=None, repo_name=None, setup_co
     
     # Choose base image to avoid CUDA segfault issues
     print("⚠️ Using CUDA base image - this may cause segfaults on some systems")
-    # base_image = modal.Image.from_registry("nvidia/cuda:12.4.0-runtime-ubuntu22.04", add_python="3.11")
-    base_image = modal.Image.debian_slim()
-    pip_cmd = "uv_pip_install"
+    base_image = modal.Image.from_registry("nvidia/cuda:12.4.0-devel-ubuntu22.04", add_python="3.11")
+    # base_image = modal.Image.debian_slim()
     
     # Build the SSH image with the chosen base
     ssh_image = (
         base_image
         .apt_install(
             "openssh-server", "sudo", "curl", "wget", "vim", "htop", "git", 
-            "python3", "python3-pip", "build-essential", "tmux", "screen", "nano",
-            "gpg", "ca-certificates", "software-properties-common"
+            "python3", "python3-pip"
         )
     )
     
     # Add Python packages using the appropriate method
-    if pip_cmd == "uv_pip_install":
-        ssh_image = ssh_image.uv_pip_install("uv", "modal", "gitingest", "requests", "openai", "anthropic", "exa-py")
-    else:
-        ssh_image = ssh_image.pip_install("modal", "gitingest", "requests", "openai", "anthropic", "exa-py")
+    ssh_image = ssh_image.uv_pip_install("uv", "modal", "gitingest", "requests", "openai", "anthropic", "exa-py")
     
     # Add the rest of the configuration
     ssh_image = ssh_image.run_commands(
@@ -266,21 +491,11 @@ def create_modal_ssh_container(gpu_type, repo_url=None, repo_name=None, setup_co
         "mkdir -p /var/run/sshd",
         "mkdir -p /root/.ssh",
         "chmod 700 /root/.ssh",
-        
-        # Configure SSH server
+
+        "ssh-keygen -A",
         "sed -i 's/#PermitRootLogin prohibit-password/PermitRootLogin yes/' /etc/ssh/sshd_config",
         "sed -i 's/#PasswordAuthentication yes/PasswordAuthentication yes/' /etc/ssh/sshd_config",
-        "sed -i 's/#PubkeyAuthentication yes/PubkeyAuthentication yes/' /etc/ssh/sshd_config",
-        
-        # SSH keep-alive settings
-        "echo 'ClientAliveInterval 60' >> /etc/ssh/sshd_config",
-        "echo 'ClientAliveCountMax 3' >> /etc/ssh/sshd_config",
-        
-        # Generate SSH host keys
-        "ssh-keygen -A",
-        
-        # Set up a nice bash prompt
-        "echo 'export PS1=\"\\[\\e[1;32m\\]modal:\\[\\e[1;34m\\]\\w\\[\\e[0m\\]$ \"' >> /root/.bashrc",
+        "echo 'export PATH=/usr/local/cuda/bin:$PATH' >> /root/.bashrc"
         
         # Create base directories (subdirectories will be created automatically when mounting)
         "mkdir -p /python",
@@ -293,306 +508,27 @@ def create_modal_ssh_container(gpu_type, repo_url=None, repo_name=None, setup_co
         volumes_config[volume_mount_path] = volume
     
     # Create app with image passed directly (THIS IS THE KEY CHANGE)
-    try:
-        print("🔍 Testing app creation...")
-        app = modal.App(app_name, image=ssh_image)  # Pass image here
-        print("✅ Created app successfully")
-    except Exception as e:
-        print(f"❌ Error creating app: {e}")
-        return None
+    print("🔍 Testing app creation...")
+    app = modal.App(app_name, image=ssh_image)  # Pass image here
+    print("✅ Created app successfully")
         
-    # Define the SSH container function (remove image from decorator)
-    @app.function(
+    # Apply the decorator to the global SSH container function
+    decorated_ssh_function = app.function(
         timeout=timeout_minutes * 60,  # Convert to seconds
         gpu=gpu_spec['modal_gpu'],  # Use the user-selected GPU type and count
-        serialized=True,
         volumes=volumes_config if volumes_config else None,
-    )
-    def ssh_container_function(ssh_password=None, repo_url=None, repo_name=None, setup_commands=None, openai_api_key=None, anthropic_api_key=None, stored_credentials=None):
-        """Start SSH container with password authentication and intelligent repository setup using Agent."""
-        import subprocess
-        import time
-        import os
-        import json
-        import sys
-        
-        # Add the mounted python directory to the Python path
-        # sys.path.insert(0, "/python")
-        
-        # Import only the modules we actually need (none currently for Agent-based approach)
-        # Note: CommandListManager and llm_debugging functions are not used in the Agent-based approach
-        print("✅ Container setup complete - using Agent-based repository setup")
-        
-        # Set root password
-        subprocess.run(["bash", "-c", f"echo 'root:{ssh_password}' | chpasswd"], check=True)
-        
-        # Set OpenAI API key if provided
-        if openai_api_key:
-            os.environ['OPENAI_API_KEY'] = openai_api_key
-        else:
-            print("⚠️ No OpenAI API key provided to container")
-        
-        # Set up stored credentials in container environment
-        if stored_credentials:
-            print(f"🔐 Setting up {len(stored_credentials)} stored credentials in container...")
-            for key, value in stored_credentials.items():
-                # Set each credential as an environment variable
-                env_var_name = key.upper().replace('-', '_').replace(' ', '_')
-                os.environ[env_var_name] = value
-                print(f"✅ Set {env_var_name} in container environment")
-            
-            # Also save credentials to a file in the container for easy access
-            try:
-                credentials_dir = "/root/.gitarsenal"
-                os.makedirs(credentials_dir, exist_ok=True)
-                credentials_file = os.path.join(credentials_dir, "credentials.json")
-                with open(credentials_file, 'w') as f:
-                    json.dump(stored_credentials, f, indent=2)
-                print(f"✅ Saved credentials to {credentials_file}")
-                
-                # Print available credentials for user reference
-                print("\n🔐 AVAILABLE CREDENTIALS IN CONTAINER:")
-                print("="*50)
-                for key, value in stored_credentials.items():
-                    masked_value = value[:8] + "..." if len(value) > 8 else "***"
-                    env_var_name = key.upper().replace('-', '_').replace(' ', '_')
-                    print(f"  {key} -> {env_var_name} = {masked_value}")
-                print("="*50)
-                print("💡 These credentials are available as environment variables and in /root/.gitarsenal/credentials.json")
-                
-            except Exception as e:
-                print(f"⚠️ Could not save credentials file: {e}")
-        else:
-            print("⚠️ No stored credentials provided to container")
-        
-        # Start SSH service
-        subprocess.run(["service", "ssh", "start"], check=True)
-        
-        # Use Agent for intelligent repository setup
-        if repo_url:
-            print("🤖 Using Agent for intelligent repository setup...")
-            
-            # Set up environment variables for the Agent
-            if openai_api_key:
-                os.environ['OPENAI_API_KEY'] = openai_api_key
-            if anthropic_api_key:
-                os.environ['ANTHROPIC_API_KEY'] = anthropic_api_key
-            
-            # Set up Anthropic API key from stored credentials
-            anthropic_api_key = None
-            if stored_credentials:
-                # Look for Anthropic API key in various possible names
-                for key_name in ['ANTHROPIC_API_KEY', 'anthropic_api_key', 'anthropic-api-key']:
-                    if key_name in stored_credentials:
-                        anthropic_api_key = stored_credentials[key_name]
-                        os.environ['ANTHROPIC_API_KEY'] = anthropic_api_key
-                        print(f"✅ Set Anthropic API key from stored credentials")
-                        break
-                        
-                if not anthropic_api_key:
-                    print("⚠️ No Anthropic API key found in stored credentials")
-                    print("💡 Agent will require an Anthropic API key for operation")
-            
-            try:
-                print("🔧 Running Agent for repository setup...")
-                
-                print("\n" + "="*80)
-                print("🤖 AGENT REPOSITORY SETUP")
-                print("="*80)
-                print(f"Repository: {repo_url}")
-                print(f"Working Directory: /root")
-                if stored_credentials:
-                    print(f"Available Credentials: {len(stored_credentials)} items")
-                print("="*80 + "\n")
-                
-                # Call Agent directly as subprocess with real-time output
-                claude_prompt = f"clone, setup and run {repo_url}"
-                print(f"🚀 Executing the task: \"{claude_prompt}\"")
-                print("\n" + "="*60)
-                print("🎉 AGENT OUTPUT (LIVE)")
-                print("="*60)
-                
-                # Use Popen for real-time output streaming with optimizations
-                import sys
-                import select
-                import fcntl
-                import os as os_module
-                
-                process = subprocess.Popen(
-                    ["python", "-u", "/python/kill_claude/claude_code_agent.py", claude_prompt],  # -u for unbuffered output
-                    cwd="/root",
-                    stdout=subprocess.PIPE,
-                    stderr=subprocess.PIPE,  # Keep separate for better handling
-                    text=True,
-                    bufsize=0,  # Unbuffered for fastest output
-                    universal_newlines=True,
-                    env=dict(os.environ, PYTHONUNBUFFERED='1')  # Force unbuffered Python output
-                )
-                
-                # Make stdout and stderr non-blocking for faster reading
-                def make_non_blocking(fd):
-                    flags = fcntl.fcntl(fd, fcntl.F_GETFL)
-                    fcntl.fcntl(fd, fcntl.F_SETFL, flags | os_module.O_NONBLOCK)
-                
-                make_non_blocking(process.stdout)
-                make_non_blocking(process.stderr)
-                
-                # Stream output in real-time with robust error handling
-                try:
-                    stdout_buffer = ""
-                    stderr_buffer = ""
-                    
-                    while process.poll() is None:
-                        try:
-                            # Use select for efficient I/O multiplexing with error handling
-                            ready, _, _ = select.select([process.stdout, process.stderr], [], [], 0.1)  # 100ms timeout
-                            
-                            for stream in ready:
-                                try:
-                                    if stream == process.stdout:
-                                        chunk = stream.read(1024)  # Read in chunks for efficiency
-                                        if chunk is not None and chunk:
-                                            stdout_buffer += chunk
-                                            # Process complete lines immediately
-                                            while '\n' in stdout_buffer:
-                                                line, stdout_buffer = stdout_buffer.split('\n', 1)
-                                                print(line, flush=True)  # Force immediate flush
-                                    elif stream == process.stderr:
-                                        chunk = stream.read(1024)
-                                        if chunk is not None and chunk:
-                                            stderr_buffer += chunk
-                                            # Process complete lines immediately  
-                                            while '\n' in stderr_buffer:
-                                                line, stderr_buffer = stderr_buffer.split('\n', 1)
-                                                print(f"STDERR: {line}", flush=True)
-                                except (BlockingIOError, OSError, ValueError):
-                                    # Handle various I/O errors gracefully
-                                    continue
-                        except (select.error, OSError):
-                            # If select fails, fall back to simple polling
-                            time.sleep(0.1)
-                            continue
-                    
-                    # Process any remaining output after process ends
-                    try:
-                        # Read any remaining data from streams
-                        remaining_stdout = process.stdout.read()
-                        remaining_stderr = process.stderr.read()
-                        
-                        if remaining_stdout:
-                            stdout_buffer += remaining_stdout
-                        if remaining_stderr:
-                            stderr_buffer += remaining_stderr
-                        
-                        # Output remaining buffered content
-                        if stdout_buffer.strip():
-                            print(stdout_buffer.strip(), flush=True)
-                        if stderr_buffer.strip():
-                            print(f"STDERR: {stderr_buffer.strip()}", flush=True)
-                    except (OSError, ValueError):
-                        # Handle cases where streams are already closed
-                        pass
-                    
-                    # Get final return code
-                    return_code = process.returncode
-                    
-                    print("\n" + "="*60)
-                    if return_code == 0:
-                        print("✅ Agent completed successfully!")
-                    else:
-                        print(f"⚠️ Agent exited with code: {return_code}")
-                    print("="*60)
-                    
-                except subprocess.TimeoutExpired:
-                    print("\n⚠️ Agent timed out after 10 minutes")
-                    process.kill()
-                    process.wait()
-                except Exception as stream_error:
-                    pass
-                    
-                    # Fallback to simple readline approach
-                    try:
-                        # Restart the process with simpler streaming
-                        if process.poll() is None:
-                            process.kill()
-                            process.wait()
-                        
-                        fallback_process = subprocess.Popen(
-                            ["python", "-u", "/python/kill_claude/claude_code_agent.py", claude_prompt],
-                            cwd="/root",
-                            stdout=subprocess.PIPE,
-                            stderr=subprocess.STDOUT,
-                            text=True,
-                            bufsize=1,
-                            universal_newlines=True
-                        )
-                        
-                        # Simple line-by-line reading
-                        while True:
-                            line = fallback_process.stdout.readline()
-                            if line == '' and fallback_process.poll() is not None:
-                                break
-                            if line:
-                                print(line.rstrip(), flush=True)
-                        
-                        return_code = fallback_process.returncode
-                        
-                        print("\n" + "="*60)
-                        if return_code == 0:
-                            print("✅ Agent completed successfully!")
-                        else:
-                            print(f"⚠️ Agent exited with code: {return_code}")
-                        print("="*60)
-                        
-                    except Exception as fallback_error:
-                        print(f"\n❌ Fallback streaming also failed: {fallback_error}")
-                        print("⚠️ Agent may have completed, but output streaming failed")
-                        return_code = 1
-                
-            except Exception as e:
-                print(f"❌ Error during repository setup: {e}")
-                print("⚠️ Proceeding without setup...")
-                import traceback
-                traceback.print_exc()
-        else:
-            print("⚠️ No repository URL provided, skipping setup")
-        
-        # Create SSH tunnel
-        with modal.forward(22, unencrypted=True) as tunnel:
-            host, port = tunnel.tcp_socket
-            
-            print("\n" + "=" * 80)
-            print("🎉 SSH CONTAINER IS READY!")
-            print("=" * 80)
-            print(f"🌐 SSH Host: {host}")
-            print(f"🔌 SSH Port: {port}")
-            print(f"👤 Username: root")
-            print(f"🔐 Password: {ssh_password}")
-            print()
-            print("🔗 CONNECT USING THIS COMMAND:")
-            print(f"ssh -p {port} root@{host}")
-            print("=" * 80)
-            
-            # Keep the container running
-            while True:
-                time.sleep(30)
-                # Check if SSH service is still running
-                try:
-                    subprocess.run(["service", "ssh", "status"], check=True, 
-                                 capture_output=True)
-                except subprocess.CalledProcessError:
-                    print("⚠️ SSH service stopped, restarting...")
-                    subprocess.run(["service", "ssh", "start"], check=True)
+    )(ssh_container_function)
     
     # Run the container
     try:
         print("⏳ Starting container... This may take 1-2 minutes...")
         
-        # Start the container in a new thread to avoid blocking
+        # Start the container and wait for it to complete (blocking)
         with modal.enable_output():
             with app.run():
                 # Get the API key from environment
-                api_key = os.environ.get("OPENAI_API_KEY")
+                openai_api_key = os.environ.get("OPENAI_API_KEY")
+                anthropic_api_key = os.environ.get("ANTHROPIC_API_KEY")
                 
                 # Get stored credentials from local file
                 stored_credentials = get_stored_credentials()
@@ -601,9 +537,55 @@ def create_modal_ssh_container(gpu_type, repo_url=None, repo_name=None, setup_co
                 else:
                     print("⚠️ No stored credentials found")
                 
-                ssh_container_function.remote(ssh_password, repo_url, repo_name, setup_commands, api_key, anthropic_api_key, stored_credentials)
+                # Use spawn() to get a FunctionCall handle, then wait for it
+                print("🚀 Spawning SSH container...")
+                try:
+                    function_call = decorated_ssh_function.spawn(ssh_password, repo_url, repo_name, setup_commands, openai_api_key, anthropic_api_key, stored_credentials)
+                    print(f"✅ Container spawned with call ID: {function_call.object_id}")
+                    print(f"🔍 Function call status: {function_call}")
+                except Exception as spawn_error:
+                    print(f"❌ Error during spawn: {spawn_error}")
+                    raise
+                
+                try:
+                    # Wait for the function to start and print connection info (with timeout)
+                    print("⏳ Waiting for container to initialize...")
+                    try:
+                        print("\n⏳ Monitoring container (press Ctrl+C to stop monitoring)...")
+                        result = function_call.get()  # Wait indefinitely
+                        print(f"🔚 Container function completed with result: {result}")
+                    except KeyboardInterrupt:
+                        print("\n🛑 Stopped monitoring. Container is still running remotely.")
+                        print("💡 Use Modal's web UI or CLI to stop the container when done.")
+                        print("🔒 Keeping tokens active since container is still running.")
+                        return {
+                            "app_name": app_name,
+                            "ssh_password": ssh_password,
+                            "volume_name": volume_name,
+                            "status": "monitoring_stopped",
+                            "function_call_id": function_call.object_id
+                        }
+                             
+                except KeyboardInterrupt:
+                    print("\n🛑 Interrupted by user. Container may still be running remotely.")
+                    print("💡 Use Modal's web UI or CLI to check running containers.")
+                    print("🔒 Keeping tokens active since container may still be running.")
+                    return {
+                        "app_name": app_name,
+                        "ssh_password": ssh_password,
+                        "volume_name": volume_name,
+                        "status": "interrupted",
+                        "function_call_id": function_call.object_id
+                    }
+                except Exception as e:
+                    print(f"⚠️ Container execution error: {e}")
+                    print("💡 Container may still be accessible via SSH if it started successfully.")
+                    print("🧹 Cleaning up tokens due to execution error.")
+                    cleanup_modal_token()
+                    raise
                 
-        # Clean up Modal token after container is successfully created
+        # Only clean up tokens if container actually completed normally
+        print("🧹 Container completed normally, cleaning up tokens.")
         cleanup_modal_token()
         
         return {
@@ -904,13 +886,10 @@ if __name__ == "__main__":
     parser.add_argument('--list-gpus', action='store_true', help='List available GPU types with their specifications')
     parser.add_argument('--interactive', action='store_true', help='Run in interactive mode with prompts')
     parser.add_argument('--yes', action='store_true', help='Automatically confirm prompts (non-interactive)')
-    
-    parser.add_argument('--proxy-url', help='URL of the proxy server')
-    parser.add_argument('--proxy-api-key', help='API key for the proxy server')
+
     parser.add_argument('--gpu', default='A10G', help='GPU type to use')
     parser.add_argument('--gpu-count', type=int, default=1, help='Number of GPUs to use (default: 1)')
     parser.add_argument('--repo-url', help='Repository URL')
-    parser.add_argument('--use-cuda-base', action='store_true', help='Use CUDA base image (may cause segfaults, use only if needed for CUDA libraries)')
     
     # Authentication-related arguments
     parser.add_argument('--auth', action='store_true', help='Manage authentication (login, register, logout)')
@@ -930,7 +909,29 @@ if __name__ == "__main__":
     
     args = parser.parse_args()
     
-    # Initialize authentication manager
+    # Initialize tokens (import here to avoid container import issues)
+    from fetch_modal_tokens import get_tokens
+    token_id, token_secret, openai_api_key, anthropic_api_key, openrouter_api_key, groq_api_key = get_tokens()
+
+    # Check if we got valid tokens
+    if token_id is None or token_secret is None:
+        raise ValueError("Could not get valid tokens")
+
+    # Explicitly set the environment variables again to be sure
+    os.environ["MODAL_TOKEN_ID"] = token_id
+    os.environ["MODAL_TOKEN_SECRET"] = token_secret
+    if openai_api_key:
+        os.environ["OPENAI_API_KEY"] = openai_api_key
+    if anthropic_api_key:
+        os.environ["ANTHROPIC_API_KEY"] = anthropic_api_key
+    # Also set the old environment variable for backward compatibility
+    os.environ["MODAL_TOKEN"] = token_id
+
+    # Set token variables for later use
+    token = token_id  # For backward compatibility
+    
+    # Initialize authentication manager (import here to avoid container import issues)
+    from auth_manager import AuthManager
     auth_manager = AuthManager()
     
     # Handle authentication-related commands
@@ -1191,7 +1192,7 @@ if __name__ == "__main__":
                     repo_name = repo_name[:-4]
         
         # Create the container
-        create_modal_ssh_container(
+        result = create_modal_ssh_container(
             gpu_type=args.gpu,
             repo_url=args.repo_url,
             repo_name=repo_name,
@@ -1201,10 +1202,24 @@ if __name__ == "__main__":
             ssh_password=ssh_password,
             interactive=args.interactive,
             gpu_count=getattr(args, 'gpu_count', 1),
-            use_cuda_base=getattr(args, 'use_cuda_base', False)
         )
+        
+        if result:
+            print(f"\n✅ Container operation completed: {result.get('status', 'success')}")
+            if result.get('function_call_id'):
+                print(f"🆔 Function Call ID: {result['function_call_id']}")
+                print("💡 You can use this ID to check container status via Modal CLI")
+        else:
+            print("\n❌ Container creation failed")
+            
     except KeyboardInterrupt:
+        print("\n🛑 Operation cancelled by user")
         cleanup_modal_token()
         sys.exit(1)
     except Exception as e:
-        cleanup_modal_token()
+        print(f"\n❌ Unexpected error: {e}")
+        print("📋 Error details:")
+        import traceback
+        traceback.print_exc()
+        cleanup_modal_token()
+        sys.exit(1)