grasp-sdk 0.1.8__py3-none-any.whl → 0.2.0a1__py3-none-any.whl

examples/test_terminal_updates.py

@@ -0,0 +1,196 @@
+#!/usr/bin/env python3
+"""
+Test script to verify terminal service updates in Python SDK.
+
+This script tests the updated terminal service functionality,
+including the new kill() method and response.json() compatibility.
+"""
+
+import asyncio
+import sys
+import os
+from unittest.mock import Mock, AsyncMock, patch
+
+# Add the parent directory to the path to import grasp_sdk
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))
+
+from grasp_sdk.services.terminal import TerminalService, StreamableCommandResult
+from grasp_sdk.services.sandbox import SandboxService, CommandEventEmitter
+from grasp_sdk.services.browser import CDPConnection
+from grasp_sdk.models import ISandboxConfig
+
+
+class MockCommandEventEmitter(CommandEventEmitter):
+    """Mock command event emitter for testing."""
+    
+    def __init__(self):
+        super().__init__()
+        self.killed = False
+        self.stdout_data = "Test output\n"
+        self.stderr_data = "Test error\n"
+        
+    async def wait(self):
+        """Mock wait method."""
+        await asyncio.sleep(0.1)  # Simulate some processing time
+        return {
+            'exit_code': 0,
+            'stdout': self.stdout_data,
+            'stderr': self.stderr_data
+        }
+    
+    async def kill(self):
+        """Mock kill method."""
+        self.killed = True
+        print("Command killed successfully")
+
+
+async def test_streamable_command_result():
+    """Test StreamableCommandResult functionality."""
+    print("\n=== Testing StreamableCommandResult ===")
+    
+    # Create mock emitter and task
+    emitter = MockCommandEventEmitter()
+    task = asyncio.create_task(emitter.wait())
+    
+    # Create StreamableCommandResult
+    result = StreamableCommandResult(emitter, task)
+    
+    # Test stdout/stderr streams
+    print("✓ StreamableCommandResult created successfully")
+    
+    # Test response.json() method
+    response_data = await result.json()
+    print(f"✓ Response JSON: {response_data}")
+    
+    # Test end() method
+    end_result = await result.end()
+    print(f"✓ End result: {end_result}")
+    
+    # Test kill() method
+    await result.kill()
+    print(f"✓ Kill method executed, emitter killed: {emitter.killed}")
+    
+    print("✓ All StreamableCommandResult tests passed!")
+
+
+async def test_terminal_service():
+    """Test TerminalService functionality."""
+    print("\n=== Testing TerminalService ===")
+    
+    # Create mock sandbox and connection
+    sandbox_config: ISandboxConfig = {
+        'key': 'test-key',
+        'timeout': 30000,
+        'debug': True
+    }
+    
+    mock_sandbox = Mock(spec=SandboxService)
+    mock_sandbox.is_debug = True
+    mock_sandbox.run_command = AsyncMock()
+    
+    mock_connection = Mock(spec=CDPConnection)
+    mock_connection.ws_url = 'ws://localhost:8080'
+    
+    # Create terminal service
+    terminal = TerminalService(mock_sandbox, mock_connection)
+    print("✓ TerminalService created successfully")
+    
+    # Mock websockets.connect
+    with patch('grasp_sdk.services.terminal.websockets.connect') as mock_connect:
+        mock_ws = AsyncMock()
+        mock_ws.open = True
+        mock_connect.return_value.__aenter__ = AsyncMock(return_value=mock_ws)
+        mock_connect.return_value.__aexit__ = AsyncMock(return_value=None)
+        # For direct await usage
+        async def mock_connect_func(*args, **kwargs):
+            return mock_ws
+        mock_connect.side_effect = mock_connect_func
+        
+        # Mock emitter
+        mock_emitter = MockCommandEventEmitter()
+        mock_sandbox.run_command.return_value = mock_emitter
+        
+        # Test run_command
+        result = await terminal.run_command('echo "Hello World"', {'timeout': 5000})
+        print("✓ run_command executed successfully")
+        
+        # Verify the command was called with correct options
+        mock_sandbox.run_command.assert_called_once()
+        call_args = mock_sandbox.run_command.call_args
+        command_arg = call_args[0][0]
+        options_arg = call_args[0][1]
+        
+        print(f"✓ Command: {command_arg}")
+        print(f"✓ Options: {options_arg}")
+        
+        # Verify inBackground option is set correctly
+        assert options_arg['inBackground'] == True, "inBackground option should be True"
+        assert options_arg['nohup'] == False, "Nohup option should be False"
+        print("✓ Command options verified")
+        
+        # Test result methods
+        response_data = await result.json()
+        print(f"✓ Response JSON: {response_data}")
+        
+        # Test kill method
+        await result.kill()
+        print("✓ Kill method executed successfully")
+    
+    # Test close method
+    terminal.close()
+    print("✓ Terminal service closed successfully")
+    
+    print("✓ All TerminalService tests passed!")
+
+
+async def test_command_options_compatibility():
+    """Test command options compatibility with TypeScript version."""
+    print("\n=== Testing Command Options Compatibility ===")
+    
+    from grasp_sdk.models import ICommandOptions
+    
+    # Test that ICommandOptions supports 'inBackground' parameter
+    options: ICommandOptions = {
+        'inBackground': True,
+        'timeout': 5000,
+        'cwd': '/tmp',
+        'nohup': False
+    }
+    
+    print(f"✓ Command options created: {options}")
+    print("✓ 'inBackground' parameter is supported in ICommandOptions")
+    
+    # Test that we can access the inBackground option
+    background_value = options.get('inBackground', False)
+    assert background_value == True, "inBackground option should be accessible"
+    print("✓ inBackground option accessible and correct")
+    
+    print("✓ All command options compatibility tests passed!")
+
+
+async def main():
+    """Run all tests."""
+    print("Starting Terminal Service Update Tests...")
+    
+    try:
+        await test_streamable_command_result()
+        await test_terminal_service()
+        await test_command_options_compatibility()
+        
+        print("\n🎉 All tests passed! Terminal service updates are working correctly.")
+        print("\n📋 Summary of updates:")
+        print("   • Added kill() method to StreamableCommandResult")
+        print("   • Added response.json() compatibility method")
+        print("   • Maintained ICommandOptions to use 'inBackground' for Python compatibility")
+        print("   • Improved error handling and WebSocket management")
+        print("   • Synchronized timeout values with TypeScript version")
+        
+    except Exception as e:
+        print(f"\n❌ Test failed: {e}")
+        import traceback
+        traceback.print_exc()
+        sys.exit(1)
+
+
+if __name__ == '__main__':
+    asyncio.run(main())

grasp_sdk/__init__.py

@@ -1,291 +1,183 @@
-"""Grasp E2B Python SDK
+"""Grasp Server Python SDK
 
-A Python SDK for E2B platform providing secure command execution 
+A Python SDK for Grasp platform providing secure command execution 
 and browser automation in isolated cloud environments.
 """
 
 import asyncio
 import signal
-import sys
-from typing import Dict, Optional, Any, Literal
+from typing import Dict, Optional, Any
+import warnings
 
 # Import utilities
-from .utils.logger import init_logger, get_logger
-from .utils.config import get_config
-from .services.browser import BrowserService
-from .services.sandbox import SandboxService
-
-# Import models and types
-from .models import (
-    ISandboxConfig,
-    IBrowserConfig,
-    ICommandOptions,
-    IScriptOptions,
-    SandboxStatus,
+from .utils.logger import get_logger
+
+# Import all classes from grasp module
+from .grasp import (
+    GraspServer,
+    GraspBrowser,
+    GraspTerminal,
+    GraspSession,
+    Grasp,
+    _servers,
+    shutdown,
 )
 
-__version__ = "0.1.8"
+# Import CDPConnection for type annotation
+from .services.browser import CDPConnection
+
+__version__ = "0.2.0a1"
 __author__ = "Grasp Team"
 __email__ = "team@grasp.dev"
 
+# Global server registry (re-exported from grasp module)
+_servers: Dict[str, GraspServer] = _servers
 
-class GraspServer:
-    """Main Grasp E2B class for browser automation."""
-    
-    def __init__(self, sandbox_config: Optional[Dict[str, Any]] = None):
-        """Initialize GraspServer with configuration.
-        
-        Args:
-            sandbox_config: Optional sandbox configuration overrides
-        """
-        if sandbox_config is None:
-            sandbox_config = {}
-        
-        # Extract browser-specific options
-        browser_type = sandbox_config.pop('type', 'chromium')
-        headless = sandbox_config.pop('headless', True)
-        adblock = sandbox_config.pop('adblock', False)
-        logLevel = sandbox_config.pop('logLevel', '')
-        keepAliveMS = sandbox_config.pop('keepAliveMS', 0)
-
-        # Set default log level
-        if not logLevel:
-            logLevel = 'debug' if sandbox_config.get('debug', False) else 'info'
-
-        self.__browser_type = browser_type
-        
-        # Create browser task
-        self.__browser_config = {
-            'headless': headless,
-            'envs': {
-                'ADBLOCK': 'true' if adblock else 'false',
-                'KEEP_ALIVE_MS': str(keepAliveMS),
-            }
-        }
-            
-        config = get_config()
-        config['sandbox'].update(sandbox_config)
-        self.config = config
-
-        logger = config['logger']
-        logger['level'] = logLevel
-        
-        # Initialize logger first
-        init_logger(logger)
-        self.logger = get_logger().child('GraspE2B')
-        
-        self.browser_service: Optional[BrowserService] = None
-        
-        self.logger.info('GraspE2B initialized')
-    
-    async def __aenter__(self):
-        connection = await self.create_browser_task()
-    
-        # Register server
-        # if connection['id']:
-        #    _servers[connection['id']] = self
-            
-        return connection
-    
-    async def __aexit__(self, exc_type, exc, tb):
-        if self.browser_service and self.browser_service.id:
-            service_id = self.browser_service.id
-            self.logger.info(f'Closing browser service {service_id}')
-            await _servers[service_id].cleanup()
-            del _servers[service_id]
-    
-    @property
-    def sandbox(self) -> Optional[SandboxService]:
-        """Get the underlying sandbox service.
-        
-        Returns:
-            SandboxService instance or None
-        """
-        return self.browser_service.get_sandbox() if self.browser_service else None
-    
-    def get_status(self) -> Optional[SandboxStatus]:
-        """Get current sandbox status.
-        
-        Returns:
-            Sandbox status or None
-        """
-        return self.sandbox.get_status() if self.sandbox else None
-    
-    def get_sandbox_id(self) -> Optional[str]:
-        """Get sandbox ID.
-        
-        Returns:
-            Sandbox ID or None
-        """
-        return self.sandbox.get_sandbox_id() if self.sandbox else None
-    
-    async def create_browser_task(
-        self, 
-    ) -> Dict[str, Any]:
-        """Create and launch a browser task.
-        
-        Args:
-            browser_type: Type of browser to launch
-            config: Browser configuration overrides
-            
-        Returns:
-            Dictionary containing browser connection info
-            
-        Raises:
-            RuntimeError: If browser service is already initialized
-        """
-        if self.browser_service:
-            raise RuntimeError('Browser service can only be initialized once')
-        
-        config = self.__browser_config
-        browser_type = self.__browser_type
-
-        if config is None:
-            config = {}
-            
-        # Create base browser config
-        browser_config: IBrowserConfig = {
-            'headless': True,
-            'launchTimeout': 30000,
-            'args': [
-                '--disable-web-security',
-                '--disable-features=VizDisplayCompositor',
-            ],
-            'envs': {},
-        }
-        
-        # Apply user config overrides with type safety
-        if 'headless' in config:
-            browser_config['headless'] = config['headless']
-        if 'launchTimeout' in config:
-            browser_config['launchTimeout'] = config['launchTimeout']
-        if 'args' in config:
-            browser_config['args'] = config['args']
-        if 'envs' in config:
-            browser_config['envs'] = config['envs']
-        
-        self.browser_service = BrowserService(
-            self.config['sandbox'],
-            browser_config
-        )
-        await self.browser_service.initialize(browser_type)
-
-        # Register server
-        _servers[str(self.browser_service.id)] = self
-        self.logger.info("🚀 Browser service initialized", {
-            'id': self.browser_service.id,
-        })
-        
-        self.logger.info('🌐 Launching Chromium browser with CDP...')
-        cdp_connection = await self.browser_service.launch_browser()
-        
-        self.logger.info('✅ Browser launched successfully!')
-        self.logger.debug(
-            f'CDP Connection Info (wsUrl: {cdp_connection.ws_url}, httpUrl: {cdp_connection.http_url})'
-        )
-        
-        return {
-            'id': self.browser_service.id,
-            'ws_url': cdp_connection.ws_url,
-            'http_url': cdp_connection.http_url
-        }
-    
-    async def cleanup(self) -> None:
-        """Cleanup resources.
-        
-        Returns:
-            Promise that resolves when cleanup is complete
-        """
-        self.logger.info('Starting cleanup process')
-        
-        try:
-            # Cleanup browser service
-            if self.browser_service:
-                await self.browser_service.cleanup()
-                
-            self.logger.info('Cleanup completed successfully')
-        except Exception as error:
-            self.logger.error(f'Cleanup failed: {error}')
-            raise
-
-
-# Global server registry
-_servers: Dict[str, GraspServer] = {}
-
-
+# Deprecated functions for backward compatibility
 async def launch_browser(
     options: Optional[Dict[str, Any]] = None
-) -> Dict[str, Any]:
+) -> CDPConnection:
     """Launch a browser instance.
     
+    .. deprecated::
+        Use grasp.launch() instead. This method will be removed in a future version.
+    
     Args:
         options: Launch options including type, headless, adblock settings
         
     Returns:
         Dictionary containing connection information
-    """    
+    """
+    import warnings
+    warnings.warn(
+        "launch_browser() is deprecated. Use grasp.launch() instead.",
+        DeprecationWarning,
+        stacklevel=2
+    )
     # Create server instance
     server = GraspServer(options)
-    
+
     connection = await server.create_browser_task()
-    
-    # Register server
-    if connection['id']:
-        _servers[connection['id']] = server
-        
+
     return connection
 
 
-async def _graceful_shutdown(signal_name: str) -> None:
-    """Handle graceful shutdown.
+async def _graceful_shutdown(reason: str = 'SIGTERM') -> None:
+    """Gracefully shutdown all servers.
     
     Args:
-        signal_name: Name of the signal received
+        reason: Reason for shutdown
     """
-    print(f'Received {signal_name}, starting cleanup...')
+    import os
+    logger = get_logger().child('shutdown')
+    logger.info(f'Graceful shutdown initiated: {reason}')
+    
+    if not _servers:
+        logger.info('No active servers to shutdown')
+        # 强制终止进程
+        os._exit(0)
+        return
     
-    # Cleanup all GraspServer instances
-    for server_id in list(_servers.keys()):
-        await _servers[server_id].cleanup()
-        del _servers[server_id]
+    logger.info(f'Shutting down {len(_servers)} active servers...')
     
-    print('All servers cleaned up, exiting...')
-    sys.exit(0)
+    # Create cleanup tasks for all servers
+    cleanup_tasks = []
+    for server_id, server in _servers.items():
+        logger.info(f'Scheduling cleanup for server {server_id}')
+        cleanup_tasks.append(server.cleanup())
+    
+    # Wait for all cleanups to complete
+    try:
+        await asyncio.gather(*cleanup_tasks, return_exceptions=True)
+        logger.info('All servers cleaned up successfully')
+    except Exception as error:
+        logger.error(f'Error during cleanup: {error}')
+    
+    # Clear the servers registry
+    _servers.clear()
+    logger.info('Graceful shutdown completed')
+    
+    # 强制终止进程，确保不会被其他异步任务阻塞
+    os._exit(0)
 
 
 def _setup_signal_handlers() -> None:
     """Setup signal handlers for graceful shutdown."""
-    def signal_handler(signum, frame):
-        signal_name = signal.Signals(signum).name
-        asyncio.create_task(_graceful_shutdown(signal_name))
-    
-    # Register signal handlers
-    signal.signal(signal.SIGINT, signal_handler)   # Ctrl+C
-    signal.signal(signal.SIGTERM, signal_handler)  # Termination signal
+    def signal_handler(signame):
+        logger = get_logger().child('signal')
+        logger.info(f'Received signal {signame}, initiating graceful shutdown...')
+        
+        # 创建关闭任务
+        asyncio.create_task(_graceful_shutdown(signame))
+    
+    # 尝试使用 asyncio 的信号处理器（更适合异步环境）
+    try:
+        loop = asyncio.get_running_loop()
+        for sig in [signal.SIGTERM, signal.SIGINT]:
+            loop.add_signal_handler(sig, lambda s=sig: signal_handler(signal.Signals(s).name))
+    except RuntimeError:
+        # 如果没有运行的事件循环，使用传统方式
+        def fallback_handler(signum, frame):
+            signal_name = signal.Signals(signum).name
+            logger = get_logger().child('signal')
+            logger.info(f'Received {signal_name}, but no event loop is running')
+            import os
+            os._exit(1)
+            
+        signal.signal(signal.SIGTERM, fallback_handler)
+        signal.signal(signal.SIGINT, fallback_handler)
 
 
-# Setup signal handlers on import
+async def shutdown(connection: Optional[str] = None) -> None:
+    """Shutdown Grasp servers (deprecated).
+    
+    Args:
+        connection: Optional connection ID to shutdown specific server
+        
+    Deprecated:
+        This function is deprecated. Use server.cleanup() or context managers instead.
+    """
+    warnings.warn(
+        "shutdown() is deprecated. Use server.cleanup() or context managers instead.",
+        DeprecationWarning,
+        stacklevel=2
+    )
+    
+    logger = get_logger().child('shutdown')
+    
+    if connection:
+        # Shutdown specific connection
+        if connection in _servers:
+            logger.info(f'Shutting down server {connection}')
+            await _servers[connection].cleanup()
+        else:
+             logger.warn(f'Server {connection} not found')
+    else:
+        # Shutdown all servers
+        await _graceful_shutdown('USER_SHUTDOWN')
+
+
+# Setup signal handlers when module is imported
 _setup_signal_handlers()
 
-
 # Export all public APIs
 __all__ = [
     'GraspServer',
+    'Grasp',
+    'GraspSession',
+    'GraspBrowser',
+    'GraspTerminal',
     'launch_browser',
-    'ISandboxConfig',
-    'IBrowserConfig',
-    'ICommandOptions', 
-    'IScriptOptions',
-    'SandboxStatus',
-    'get_config',
-    'init_logger',
-    'get_logger',
-    'BrowserService',
-    'SandboxService',
+    'shutdown',
 ]
 
-
 # Default export equivalent
 default = {
     'GraspServer': GraspServer,
+    'Grasp': Grasp,
+    'GraspSession': GraspSession,
+    'GraspBrowser': GraspBrowser,
+    'GraspTerminal': GraspTerminal,
+    'launch_browser': launch_browser,
+    'shutdown': shutdown,
 }

grasp_sdk/grasp/__init__.py

@@ -0,0 +1,26 @@
+"""Grasp SDK classes and utilities."""
+
+from typing import Dict
+
+from .server import GraspServer
+from .browser import GraspBrowser
+from .terminal import GraspTerminal
+from .session import GraspSession
+from .index import Grasp
+from .utils import (
+    _servers,
+    shutdown,
+)
+
+# Export the global servers registry
+_servers: Dict[str, GraspServer] = _servers
+
+__all__ = [
+    'GraspServer',
+    'GraspBrowser', 
+    'GraspTerminal',
+    'GraspSession',
+    'Grasp',
+    'shutdown',
+    '_servers',
+]