grasp-sdk 0.1.9__py3-none-any.whl → 0.2.0b1__py3-none-any.whl

grasp_sdk/services/sandbox.py

@@ -7,6 +7,7 @@ equivalent to the TypeScript version in src/services/sandbox.service.ts
 """
 
 import asyncio
+import json
 import os
 import tempfile
 import time
@@ -41,6 +42,8 @@ class CommandEventEmitter:
     def set_handle(self, handle: Any) -> None:
         """Set the command handle (used internally)."""
         self.handle = handle
+        if self.is_killed:
+            self.handle.kill()
     
     def on(self, event: str, callback: Callable) -> None:
         """Register event listener."""
@@ -48,6 +51,11 @@ class CommandEventEmitter:
             self._callbacks[event] = []
         self._callbacks[event].append(callback)
     
+    def off(self, event: str, callback: Callable) -> None:
+        """Remove event listener."""
+        if event in self._callbacks and callback in self._callbacks[event]:
+            self._callbacks[event].remove(callback)
+    
     def emit(self, event: Any, *args) -> None:
         """Emit event to all registered listeners."""
         # print(f'🚀 emit ${event}')
@@ -60,16 +68,41 @@ class CommandEventEmitter:
                     pass
     
     async def wait(self) -> Any:
-        """Wait for command to complete."""
+        """Wait for command to complete.
+        
+        Returns:
+            Command result or exception result if command fails
+        """
         while not self.handle:
             await asyncio.sleep(0.1)
-        return await self.handle.wait()
+        try:
+            return await self.handle.wait()
+        except Exception as ex:
+            # 对于 CommandExitException，提取有用信息并返回结构化结果
+            if hasattr(ex, 'exit_code') or 'CommandExitException' in str(type(ex)):
+                return {
+                    'error': str(ex),
+                    'exit_code': getattr(ex, 'exit_code', -1),
+                    'stdout': getattr(ex, 'stdout', ''),
+                    'stderr': getattr(ex, 'stderr', str(ex))
+                }
+            # Return exception result if available, similar to TypeScript version
+            if hasattr(ex, 'result'):
+                # 尝试获取异常的结果属性,如果不存在则返回None
+                return getattr(ex, 'result', None)
+            # 对于其他异常，返回结构化错误信息
+            return {
+                'error': str(ex),
+                'exit_code': -1,
+                'stdout': '',
+                'stderr': str(ex)
+            }
     
     async def kill(self) -> None:
         """Kill the running command."""
-        if not self.handle:
-            raise RuntimeError('Command handle not set')
-        await self.handle.kill()
+        self.is_killed = True
+        if self.handle:
+            await self.handle.kill()
     
     def get_handle(self) -> Optional[Any]:
         """Get the original command handle."""
@@ -123,6 +156,56 @@ class SandboxService:
         """Get timeout value."""
         return self.config['timeout']
     
+    async def connect_sandbox(self, sandbox_id: str) -> None:
+        """
+        Connects to an existing sandbox.
+        
+        Args:
+            sandbox_id: ID of the sandbox to connect to
+            
+        Raises:
+            RuntimeError: If sandbox connection fails
+        """
+        try:
+            self.status = SandboxStatus.CREATING
+            
+            self.logger.info(f'Connection Grasp sandbox: {sandbox_id}')
+            
+            # Verify authentication
+            res = await verify({'key': self.config['key']})
+            if not res['success']:
+                raise RuntimeError('Authorization failed.')
+            
+            api_key = res['data']['token']
+            
+            # Connect to existing sandbox
+            self.sandbox = await AsyncSandbox.connect(
+                sandbox_id=sandbox_id,
+                api_key=api_key
+            )
+            
+            self.status = SandboxStatus.RUNNING
+            
+            # Read existing config from sandbox
+            config_content = await self.sandbox.files.read('/home/user/.grasp-config.json')
+            timeout = self.config['timeout']
+            self.config = json.loads(config_content)
+            
+            # Update timeout if different
+            if timeout > 0 and timeout != self.config['timeout']:
+                # Note: e2b Python SDK may not have setTimeout method
+                # This is equivalent to the TypeScript version's setTimeout call
+                self.sandbox.set_timeout(timeout)
+            
+            self.logger.info('Grasp sandbox connected', {
+                'sandboxId': getattr(self.sandbox, 'sandbox_id', 'unknown'),
+            })
+            
+        except Exception as error:
+            self.status = SandboxStatus.ERROR
+            self.logger.error('Failed to connect Grasp sandbox', str(error))
+            raise RuntimeError(f'Failed to connect sandbox: {error}')
+    
     async def create_sandbox(self, template_id: str, envs: Optional[Dict[str, str]] = None) -> None:
         """
         Creates and starts a new sandbox.
@@ -135,7 +218,7 @@ class SandboxService:
             self.status = SandboxStatus.CREATING
             
             # Verify authentication
-            res = await verify(self.config)
+            res = await verify({'key': self.config['key']})
             if not res['success']:
                 raise RuntimeError('Authorization failed.')
             
@@ -152,6 +235,17 @@ class SandboxService:
                 timeout=self.config['timeout'] // 1000,  # Convert ms to seconds
                 envs=envs,
             )
+            
+            # Write grasp config file
+            import json
+            config_data = {
+                'id': getattr(self.sandbox, 'sandbox_id', 'unknown'),
+                **self.config
+            }
+            await self.sandbox.files.write(
+                '/home/user/.grasp-config.json',
+                json.dumps(config_data)
+            )
 
             self.status = SandboxStatus.RUNNING
             self.logger.info('Grasp sandbox created successfully', {
@@ -192,10 +286,11 @@ class SandboxService:
             options = {}
         
         cwd = options.get('cwd', self.DEFAULT_WORKING_DIRECTORY)
-        timeout_ms = options.get('timeout', self.config['timeout'])
+        timeout_ms = options.get('timeout_ms', 0)
         use_nohup = options.get('nohup', False)
         in_background = options.get('inBackground', False)
         envs = options.get('envs', {})
+        user = options.get('user', 'user')
 
         # print(f'command: {command}')
         # print(f'🚀 options {options}')
@@ -224,9 +319,10 @@ class SandboxService:
                     result = await self.sandbox.commands.run(
                         final_command,
                         cwd=cwd,
-                        timeout=timeout_ms,
+                        timeout=timeout_ms // 1000,
                         envs=envs,
                         background=in_background,
+                        user=user,
                     )
                 else:
                     raise RuntimeError('Sandbox commands interface not available')
@@ -252,6 +348,7 @@ class SandboxService:
                         timeout=timeout_ms // 1000,
                         envs=envs,
                         background=in_background,
+                        user=user,
                     )
                     return None
 
@@ -271,14 +368,9 @@ class SandboxService:
                                 timeout=timeout_ms // 1000,
                                 background=True,
                                 envs=envs,
-                                on_stdout=lambda data: (
-                                    self.logger.debug('Command stdout', {'data': data}),
-                                    event_emitter.emit('stdout', data)
-                                )[1],
-                                on_stderr=lambda data: (
-                                    self.logger.debug('Command stderr', {'data': data}),
-                                    event_emitter.emit('stderr', data)
-                                )[1]
+                                user=user,
+                                on_stdout=lambda data: event_emitter.emit('stdout', data),
+                                on_stderr=lambda data: event_emitter.emit('stderr', data)
                             )
                         
                             event_emitter.set_handle(handle)
@@ -320,11 +412,11 @@ class SandboxService:
         options: Optional[IScriptOptions] = None
     ) -> Union[Any, CommandEventEmitter]:
         """
-        Runs JavaScript code in the sandbox.
+        Runs JavaScript or Python code in the sandbox.
         
         Args:
-            code: JavaScript code to execute, or file path starting with '/home/user/'
-            options: Script execution options
+            code: JavaScript/Python code to execute, or file path starting with '/home/user/'
+            options: Script execution options (type: 'cjs', 'esm', or 'py')
             
         Returns:
             CommandResult for synchronous execution,
@@ -335,7 +427,8 @@ class SandboxService:
             
         Note:
             If code starts with '/home/user/', it will be treated as a file path.
-            Otherwise, it will be treated as JavaScript code and written to a temporary file.
+            Otherwise, it will be treated as code and written to a temporary file.
+            For Python scripts, use type='py' and the code will be executed with python3.
         """
         if not self.sandbox or self.status != SandboxStatus.RUNNING:
             raise RuntimeError('Sandbox is not running. Call create_sandbox() first.')
@@ -349,13 +442,18 @@ class SandboxService:
             script_path = code
             
             if not code.startswith('/home/user/'):
-                extension = 'mjs' if options['type'] == 'esm' else 'js'
+                if options['type'] == 'py':
+                    extension = 'py'
+                elif options['type'] == 'esm':
+                    extension = 'mjs'
+                else:
+                    extension = 'js'
                 working_dir = options.get('cwd', self.DEFAULT_WORKING_DIRECTORY)
                 script_path = f'{working_dir}/script_{timestamp}.{extension}'
                 # Write code to temporary file
                 await self.sandbox.files.write(script_path, code)
             
-            self.logger.debug('Running JavaScript code in sandbox', {
+            self.logger.debug('Running script in sandbox', {
                 'type': options['type'],
                 'scriptPath': script_path,
                 'codeLength': len(code),
@@ -363,7 +461,10 @@ class SandboxService:
             
             # Choose execution command based on type
             pre_command = options.get('preCommand', '')
-            command = f'{pre_command}node {script_path}'
+            if options['type'] == 'py':
+                command = f'{pre_command}python3 {script_path}'
+            else:
+                command = f'{pre_command}node {script_path}'
             
             # Set environment variables
             envs = options.get('envs', {})
@@ -372,10 +473,11 @@ class SandboxService:
             # Prepare command options
             cmd_options: Optional[Any] = {
                 'cwd': options.get('cwd'),
-                'timeout': options.get('timeoutMs'),
+                'timeout_ms': options.get('timeout_ms', 0),
                 'inBackground': options.get('background', False),
                 'nohup': options.get('nohup', False),
                 'envs': envs,
+                'user': options.get('user', 'user'),
             }
             
             # Execute the script
@@ -408,6 +510,86 @@ class SandboxService:
         ext = Path(file_path).suffix.lower()
         return ext in binary_extensions
     
+    def _encode_content(self, data: bytes, encoding: str) -> Union[str, bytes]:
+        """
+        Encode content based on specified encoding.
+        
+        Args:
+            data: Raw bytes data
+            encoding: Encoding type ('utf8', 'base64', or 'binary')
+            
+        Returns:
+            Encoded content as string or bytes
+        """
+        if encoding == 'utf8':
+            return data.decode('utf-8')
+        elif encoding == 'base64':
+            import base64
+            return base64.b64encode(data).decode('utf-8')
+        else:  # binary
+            return data
+    
+    async def read_file_from_sandbox(
+        self, 
+        remote_path: str, 
+        options: Optional[Dict[str, str]] = None
+    ) -> Union[str, bytes]:
+        """
+        Read content from a file in the sandbox.
+        
+        Args:
+            remote_path: File path in sandbox to read from
+            options: Dictionary with 'encoding' key ('utf8', 'base64', or 'binary')
+            
+        Returns:
+            File content as string or bytes depending on encoding
+            
+        Raises:
+            RuntimeError: If sandbox is not running or file read fails
+        """
+        if not self.sandbox or self.status != SandboxStatus.RUNNING:
+            raise RuntimeError('Sandbox is not running. Call create_sandbox() first.')
+        
+        if options is None:
+            options = {}
+        
+        # Determine encoding
+        encoding = options.get('encoding')
+        if not encoding:
+            encoding = 'binary' if self._is_binary_file(remote_path) else 'utf8'
+        
+        try:
+            self.logger.debug('Reading file from sandbox', {
+                'remotePath': remote_path,
+                'encoding': encoding
+            })
+            
+            # Always read as bytes first
+            file_content = await self.sandbox.files.read(remote_path, format='bytes')
+            
+            # Encode based on requested format
+            return self._encode_content(file_content, encoding)
+            
+        except Exception as error:
+            self.logger.error(f'Failed to read file from sandbox: {error}')
+            raise RuntimeError(f'Failed to read file: {error}')
+    
+    async def write_file_to_sandbox(self, remote_path: str, content: Union[str, bytes]) -> None:
+        """
+        Write content directly to a file in the sandbox.
+        
+        Args:
+            remote_path: File path in sandbox where content will be written
+            content: String or bytes content to write to the file
+            
+        Raises:
+            RuntimeError: If sandbox is not running
+        """
+        if not self.sandbox or self.status != SandboxStatus.RUNNING:
+            raise RuntimeError('Sandbox is not running. Call create_sandbox() first.')
+        
+        return await self.sandbox.files.write(remote_path, content)
+    
     async def copy_file_from_sandbox(self, remote_path: str, local_path: str) -> None:
         """
         Copy file from sandbox to local filesystem.
@@ -442,7 +624,7 @@ class SandboxService:
                 with open(local_path, 'w', encoding='utf-8') as f:
                     f.write(file_content)
             
-            self.logger.info(f'File copied from sandbox: {remote_path} -> {local_path}')
+            self.logger.debug(f'File copied from sandbox: {remote_path} -> {local_path}')
         
         except Exception as error:
             self.logger.error(f'Failed to copy file: {error}')
@@ -529,6 +711,201 @@ class SandboxService:
             self.logger.error('Failed to list files', str(error))
             raise RuntimeError(f'Failed to list files: {error}')
     
+    def _validate_safe_path(self, path: str) -> bool:
+        """
+        Validate if a path is safe for synchronization.
+        
+        Args:
+            path: Path to validate
+            
+        Returns:
+            True if path is safe, False otherwise
+        """
+        # Normalize path
+        normalized_path = path.replace('//', '/').rstrip('/')
+        
+        # Check if it's an absolute path under /home/user
+        if normalized_path.startswith('/'):
+            # Absolute path must be under /home/user/ and not /home/user itself
+            return (
+                normalized_path.startswith('/home/user/') and
+                normalized_path != '/home/user'
+            )
+        
+        # Relative path checks
+        # Don't allow .. to access parent directories
+        if '..' in normalized_path:
+            return False
+        
+        # Don't allow syncing current directory itself (. or empty string)
+        if normalized_path in ('.', '', './'):
+            return False
+        
+        # Relative path must be a subdirectory (starts with ./ or doesn't start with /)
+        return normalized_path.startswith('./') or not normalized_path.startswith('/')
+    
+    async def _is_directory(self, remote_path: str) -> bool:
+        """
+        Check if a remote path is a directory.
+        
+        Args:
+            remote_path: Remote path to check
+            
+        Returns:
+            True if path is a directory, False otherwise
+        """
+        if not self.sandbox or self.status != SandboxStatus.RUNNING:
+            return False
+        
+        try:
+            result = await self.run_command(f'test -d "{remote_path}" && echo "true" || echo "false"')
+            stdout = getattr(result, 'stdout', '').strip() if result else ''
+            return stdout == 'true'
+        except Exception as error:
+            self.logger.debug('Error checking if path is directory', {
+                'remotePath': remote_path,
+                'error': str(error)
+            })
+            return False
+    
+    async def _sync_directory_recursive(
+        self, 
+        remote_path: str, 
+        local_path: str
+    ) -> int:
+        """
+        Recursively sync a directory from sandbox to local filesystem.
+        
+        Args:
+            remote_path: Remote directory path
+            local_path: Local directory path
+            
+        Returns:
+            Number of files synced
+        """
+        synced_count = 0
+        
+        try:
+            # Get all items in the directory
+            items = await self.list_files(remote_path)
+            
+            if len(items) == 0:
+                self.logger.debug('No items found in directory', {'remotePath': remote_path})
+                return 0
+            
+            # Ensure local directory exists
+            os.makedirs(local_path, exist_ok=True)
+            
+            # Process each item
+            for item in items:
+                remote_item_path = f"{remote_path}/{item}"
+                local_item_path = os.path.join(local_path, item)
+                
+                # Check if it's a directory
+                is_dir = await self._is_directory(remote_item_path)
+                
+                if is_dir:
+                    self.logger.debug('Syncing subdirectory', {
+                        'from': remote_item_path,
+                        'to': local_item_path
+                    })
+                    
+                    # Recursively sync subdirectory
+                    sub_synced_count = await self._sync_directory_recursive(
+                        remote_item_path,
+                        local_item_path
+                    )
+                    synced_count += sub_synced_count
+                else:
+                    self.logger.debug('Syncing file', {
+                        'from': remote_item_path,
+                        'to': local_item_path
+                    })
+                    
+                    # Sync file
+                    await self.copy_file_from_sandbox(remote_item_path, local_item_path)
+                    synced_count += 1
+            
+            return synced_count
+        
+        except Exception as error:
+            self.logger.error('Error in recursive directory sync', {
+                'remotePath': remote_path,
+                'localPath': local_path,
+                'error': str(error)
+            })
+            raise
+    
+    async def sync_downloads_directory(
+        self, 
+        dist: str = '/tmp/grasp/downloads', 
+        src: str = './downloads'
+    ) -> str:
+        """
+        Synchronize sandbox downloads directory to local filesystem (recursive sync of all subdirectories).
+        
+        Args:
+            dist: Local target directory path, defaults to '/tmp/grasp/downloads'
+            src: Remote source directory path, defaults to './downloads'
+            
+        Returns:
+            Local sync directory path
+            
+        Raises:
+            RuntimeError: If sandbox is not running or sync fails
+            ValueError: If path is unsafe
+        """
+        # Validate path safety
+        if not self._validate_safe_path(src):
+            error_msg = f"Unsafe path detected: {src}. Only subdirectories under /home/user are allowed."
+            self.logger.error('Path validation failed', {
+                'path': src,
+                'error': error_msg
+            })
+            raise ValueError(error_msg)
+        
+        remote_path = src
+        local_path = dist
+        
+        try:
+            self.logger.debug('🗂️ Starting recursive directory sync', {
+                'remotePath': remote_path,
+                'localPath': local_path
+            })
+            
+            # Check if remote directory exists
+            dir_exists = await self._is_directory(remote_path)
+            if not dir_exists:
+                # Try to handle as file
+                parent_dir = os.path.dirname(remote_path)
+                file_name = os.path.basename(remote_path)
+                
+                files = await self.list_files(parent_dir)
+                
+                if file_name not in files:
+                    self.logger.debug('Remote path does not exist', {'remotePath': remote_path})
+                    return ''
+            
+            # Recursively sync directory
+            synced_count = await self._sync_directory_recursive(
+                remote_path,
+                local_path
+            )
+            
+            if synced_count == 0:
+                self.logger.debug('No files found to sync')
+                return ''
+            
+            self.logger.info(
+                f'🎉 Successfully synced {synced_count} files recursively from {remote_path}'
+            )
+            
+            return local_path
+        
+        except Exception as error:
+            self.logger.error('Error syncing downloads directory:', str(error))
+            raise
+    
     def get_status(self) -> SandboxStatus:
         """Gets the current sandbox status."""
         return self.status
@@ -572,7 +949,8 @@ class SandboxService:
                     'sandboxId': getattr(self.sandbox, 'sandbox_id', 'unknown'),
                 })
                 
-                await self.sandbox.kill()
+                if await self.sandbox.is_running():
+                    await self.sandbox.kill()
                 self.sandbox = None
                 self.status = SandboxStatus.STOPPED