refactor: remove AWSVMManagerWithProxy and integrate proxy support directly into AWSVMManager for streamlined VM allocation;

minor fix on openai_cua_agent

desktop_env/providers/aws/manager.py

@@ -18,11 +18,16 @@ if not os.getenv('AWS_SUBNET_ID') or not os.getenv('AWS_SECURITY_GROUP_ID'):
 
 from desktop_env.providers.base import VMManager
 
+# Import proxy-related modules only when needed
+try:
+    from desktop_env.providers.aws.proxy_pool import get_global_proxy_pool, init_proxy_pool
+    PROXY_SUPPORT_AVAILABLE = True
+except ImportError:
+    PROXY_SUPPORT_AVAILABLE = False
+
 logger = logging.getLogger("desktopenv.providers.aws.AWSVMManager")
 logger.setLevel(logging.INFO)
 
-REGISTRY_PATH = '.aws_vms'
-
 DEFAULT_REGION = "us-east-1"
 # todo: Add doc for the configuration of image, security group and network interface
 # todo: public the AMI images
@@ -118,17 +123,55 @@ def _allocate_vm(region=DEFAULT_REGION):
     return instance_id
 
 
+def _allocate_vm_with_proxy(region=DEFAULT_REGION, proxy_config_file=None):
+    """Allocate a VM with proxy configuration"""
+    if not PROXY_SUPPORT_AVAILABLE:
+        logger.warning("Proxy support not available, falling back to regular VM allocation")
+        return _allocate_vm(region)
+    
+    from desktop_env.providers.aws.provider_with_proxy import AWSProviderWithProxy
+    
+    # Initialize proxy pool if needed
+    if proxy_config_file:
+        init_proxy_pool(proxy_config_file)
+    
+    # Get current proxy
+    proxy_pool = get_global_proxy_pool()
+    current_proxy = proxy_pool.get_next_proxy()
+    
+    if current_proxy:
+        logger.info(f"Allocating VM with proxy: {current_proxy.host}:{current_proxy.port}")
+    
+    # Create provider instance
+    provider = AWSProviderWithProxy(region=region, proxy_config_file=proxy_config_file)
+    
+    # Create new instance
+    instance_id = provider.create_instance_with_proxy(
+        image_id=IMAGE_ID_MAP[region],
+        instance_type=INSTANCE_TYPE,
+        security_groups=[os.getenv('AWS_SECURITY_GROUP_ID')],
+        subnet_id=os.getenv('AWS_SUBNET_ID')
+    )
+    
+    return instance_id
+
+
 class AWSVMManager(VMManager):
     """
     AWS VM Manager for managing virtual machines on AWS.
     
     AWS does not need to maintain a registry of VMs, as it can dynamically allocate and deallocate VMs.
-    This class remains the interface of VMManager for compatibility with other components.
+    This class supports both regular VM allocation and proxy-enabled VM allocation.
     """
-    def __init__(self, registry_path=REGISTRY_PATH):
-        self.registry_path = registry_path
+    def __init__(self, proxy_config_file=None, **kwargs):
+        self.proxy_config_file = proxy_config_file
         # self.lock = FileLock(".aws_lck", timeout=60)
         self.initialize_registry()
+        
+        # Initialize proxy pool if proxy configuration is provided
+        if proxy_config_file and PROXY_SUPPORT_AVAILABLE:
+            init_proxy_pool(proxy_config_file)
+            logger.info(f"Proxy pool initialized with config: {proxy_config_file}")
 
     def initialize_registry(self, **kwargs):
         pass
@@ -164,6 +207,10 @@ class AWSVMManager(VMManager):
         pass
 
     def get_vm_path(self, region=DEFAULT_REGION, **kwargs):
-        logger.info("Allocating a new VM in region: {}".format(region))
-        new_vm_path = _allocate_vm(region)
+        if self.proxy_config_file:
+            logger.info("Allocating a new VM with proxy configuration in region: {}".format(region))
+            new_vm_path = _allocate_vm_with_proxy(region, self.proxy_config_file)
+        else:
+            logger.info("Allocating a new VM in region: {}".format(region))
+            new_vm_path = _allocate_vm(region)
         return new_vm_path