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

grasp_sdk/grasp/index.py

@@ -0,0 +1,127 @@
+"""Main Grasp class for SDK initialization."""
+
+import os
+from typing import Dict, Optional, Any
+
+from .server import GraspServer
+from .session import GraspSession
+from ..utils.config import get_config
+
+
+class Grasp:
+    """Main Grasp SDK class for creating and managing sessions."""
+    
+    def __init__(self, options: Optional[Dict[str, Any]] = None):
+        """Initialize Grasp SDK.
+        
+        Args:
+            options: Configuration options including apiKey
+        """
+        if options is None:
+            options = {}
+        
+        self.key = options.get('apiKey', os.environ.get('GRASP_KEY', ''))
+        self._session: Optional[GraspSession] = None
+        self._launch_options: Optional[Dict[str, Any]] = None
+    
+    async def launch(self, options: Dict[str, Any]) -> GraspSession:
+        """Launch a new browser session.
+        
+        Args:
+            options: Launch options containing browser configuration
+            
+        Returns:
+            GraspSession instance
+        """
+        config = get_config()
+        browser_options = options.get('browser', {})
+        browser_options['key'] = self.key or config['sandbox']['key']
+        browser_options['keepAliveMS'] = options.get('keepAliveMS', config['sandbox']['keepAliveMs'])
+        browser_options['timeout'] = options.get('timeout', config['sandbox']['timeout'])
+        browser_options['debug'] = options.get('debug', config['sandbox']['debug'])
+        browser_options['logLevel'] = options.get('logLevel', config['logger']['level'])
+        
+        server = GraspServer(browser_options)
+        cdp_connection = await server.create_browser_task()
+        return GraspSession(cdp_connection)
+    
+    async def __aenter__(self) -> GraspSession:
+        """异步上下文管理器入口，启动会话。
+        
+        Returns:
+            GraspSession 实例
+            
+        Raises:
+            RuntimeError: 如果已经有活跃的会话
+        """
+        if self._session is not None:
+            raise RuntimeError('Session already active. Close existing session before creating a new one.')
+        
+        if self._launch_options is None:
+            raise RuntimeError('No launch options provided. Use grasp.launch(options) as context manager.')
+        
+        self._session = await self.launch(self._launch_options)
+        return self._session
+    
+    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
+        """异步上下文管理器退出，清理会话。
+        
+        Args:
+            exc_type: 异常类型
+            exc_val: 异常值
+            exc_tb: 异常回溯
+        """
+        if self._session is not None:
+            await self._session.close()
+            self._session = None
+        self._launch_options = None
+    
+    def launch_context(self, options: Dict[str, Any]) -> 'Grasp':
+        """设置启动选项并返回自身，用于上下文管理器。
+        
+        Args:
+            options: 启动选项
+            
+        Returns:
+            自身实例，用于 async with 语法
+            
+        Example:
+            async with grasp.launch_context({
+                'browser': {
+                    'type': 'chromium',
+                    'headless': True,
+                    'timeout': 30000
+                }
+            }) as session:
+                # 使用 session
+                pass
+        """
+        self._launch_options = options
+        return self
+    
+    async def connect(self, session_id: str) -> GraspSession:
+        """Connect to an existing session.
+        
+        Args:
+            session_id: ID of the session to connect to
+            
+        Returns:
+            GraspSession instance
+        """
+        # Check if server already exists
+        from . import _servers
+        server = _servers.get(session_id)
+        cdp_connection = None
+        
+        if server:
+            cdp_connection = server.browser_service.get_cdp_connection() if server.browser_service else None
+        
+        if not cdp_connection:
+            # Create new server with no timeout for existing sessions
+            server = GraspServer({
+                'timeout': 0,
+                'key': self.key,
+            })
+            cdp_connection = await server.connect_browser_task(session_id)
+        
+        return GraspSession(cdp_connection)

grasp_sdk/grasp/server.py

@@ -0,0 +1,250 @@
+"""GraspServer class for browser automation."""
+
+import os
+from typing import Dict, Optional, Any
+
+from ..utils.logger import init_logger, get_logger
+from ..utils.config import get_config
+from ..services.browser import BrowserService, CDPConnection
+from ..services.sandbox import SandboxService
+from ..models import (
+    ISandboxConfig,
+    IBrowserConfig,
+    SandboxStatus,
+)
+
+
+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('GraspServer')
+        
+        self.browser_service: Optional[BrowserService] = None
+        
+        self.logger.info('GraspServer 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}')
+            from . import _servers
+            await _servers[service_id].cleanup()
+    
+    @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, 
+    ) -> CDPConnection:
+        """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
+        from . import _servers
+        _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 cdp_connection
+    
+    async def connect_browser_task(self, session_id: str) -> CDPConnection:
+        """Connect to an existing browser session.
+        
+        Args:
+            session_id: ID of the session to connect to
+            
+        Returns:
+            CDPConnection instance for the existing session
+            
+        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
+        )
+        
+        # Connect to existing session
+        connection = await self.browser_service.connect(session_id)
+        
+        # Register server
+        from . import _servers
+        _servers[connection.id] = self
+        
+        return connection
+    
+    async def cleanup(self) -> None:
+        """Cleanup resources.
+        
+        Returns:
+            Promise that resolves when cleanup is complete
+        """
+        self.logger.info('Starting cleanup process')
+        
+        try:
+            if self.browser_service:
+                id = self.browser_service.id
+                await self.browser_service.cleanup()
+                if id:
+                    from . import _servers
+                    del _servers[id]
+                
+            self.logger.info('Cleanup completed successfully')
+        except Exception as error:
+            self.logger.error(f'Cleanup failed: {error}')
+            raise

grasp_sdk/grasp/session.py

@@ -0,0 +1,146 @@
+"""GraspSession class for session management."""
+
+import asyncio
+import websockets
+from typing import Optional
+from .browser import GraspBrowser
+from ..services.terminal import TerminalService
+from ..services.browser import CDPConnection
+from ..services.filesystem import FileSystemService
+from ..utils.logger import get_logger
+
+
+class GraspSession:
+    """Main session interface providing access to browser, terminal, and files."""
+    
+    def __init__(self, connection: CDPConnection):
+        """Initialize GraspSession.
+        
+        Args:
+            connection: CDP connection instance
+        """
+        self.connection = connection
+        self.browser = GraspBrowser(connection, self)
+        
+        # Create files service
+        connection_id = connection.id
+        try:
+            from .utils import _servers
+            server = _servers[connection_id]
+            if server.sandbox is None:
+                raise RuntimeError('Sandbox is not available for file system service')
+            self.files = FileSystemService(server.sandbox, connection)
+            self.terminal = TerminalService(server.sandbox, connection)
+        except (ImportError, KeyError) as e:
+            # In test environment or when server is not available, create a mock files service
+            self.logger.debug(f"Warning: Files service not available: {e}")
+            raise e
+        
+        # WebSocket connection for keep-alive
+        self._ws = None
+        self._is_closed = False
+
+        self.logger = get_logger().child('GraspSession')
+        
+        # Initialize WebSocket connection and start keep-alive
+        # self.logger.debug('create session...')
+        asyncio.create_task(self._initialize_websocket())
+    
+    async def _initialize_websocket(self) -> None:
+        """Initialize WebSocket connection and start keep-alive."""
+        try:
+            self._ws = await websockets.connect(self.connection.ws_url)
+            if not self._is_closed:
+                await self._keep_alive()
+            else:
+                await self._ws.close()
+        except Exception as e:
+            self.logger.debug(f"Failed to initialize WebSocket connection: {e}")
+    
+    async def _keep_alive(self) -> None:
+        """Keep WebSocket connection alive with periodic ping."""
+        ws = self._ws
+        await asyncio.sleep(10)  # 10 seconds interval, same as TypeScript version
+        if ws and ws.close_code is None:
+            try:
+                # Send ping and wait for pong response
+                pong_waiter = ws.ping()
+                await pong_waiter
+                # Recursively call keep_alive after receiving pong
+                await self._keep_alive()
+            except Exception as e:
+                self.logger.debug(f"Keep-alive ping failed: {e}")
+    
+    @property
+    def id(self) -> str:
+        """Get session ID.
+        
+        Returns:
+            Session ID
+        """
+        return self.connection.id
+    
+    def is_running(self) -> bool:
+        """Check if the session is currently running.
+        
+        Returns:
+            True if the session is running, False otherwise
+        """
+        try:
+            from .utils import _servers
+            from ..models import SandboxStatus
+            server = _servers[self.id]
+            if server.sandbox is None:
+                return False
+            status = server.sandbox.get_status()
+            return status == SandboxStatus.RUNNING
+        except (ImportError, KeyError) as e:
+            self.logger.error(f"Failed to check running status: {e}")
+            return False
+    
+    def get_host(self, port: int) -> Optional[str]:
+        """Get the external host address for a specific port.
+        
+        Args:
+            port: Port number to get host for
+            
+        Returns:
+            External host address or None if sandbox not available
+        """
+        try:
+            from .utils import _servers
+            server = _servers[self.id]
+            if server.sandbox is None:
+                self.logger.warn('Cannot get host: sandbox not available')
+                return None
+            return server.sandbox.get_sandbox_host(port)
+        except (ImportError, KeyError) as e:
+            self.logger.error(f"Failed to get host: {e}")
+            return None
+    
+    async def close(self) -> None:
+        """Close the session and cleanup resources."""
+        # Set closed flag to prevent new keep-alive cycles
+        self._is_closed = True
+        
+        # Close WebSocket connection with timeout
+        if self._ws and self._ws.close_code is None:
+            try:
+                # Add timeout to prevent hanging on close
+                await asyncio.wait_for(self._ws.close(), timeout=5.0)
+                self.logger.debug('✅ WebSocket closed successfully')
+            except asyncio.TimeoutError:
+                self.logger.debug('⚠️ WebSocket close timeout, forcing termination')
+                # Force close if timeout
+                if hasattr(self._ws, 'transport') and self._ws.transport:
+                    self._ws.transport.close()
+            except Exception as e:
+                self.logger.error(f"Failed to close WebSocket connection: {e}")
+        
+        # Cleanup server resources
+        try:
+            from .utils import _servers
+            await _servers[self.id].cleanup()
+        except (ImportError, KeyError) as e:
+            # In test environment or when server is not available, skip cleanup
+            self.logger.error(f"Warning: Server cleanup not available: {e}")

grasp_sdk/grasp/utils.py

@@ -0,0 +1,90 @@
+"""Utility functions for Grasp SDK."""
+
+import os
+import asyncio
+import signal
+from typing import Dict, Optional, List
+import warnings
+
+from ..utils.logger import get_logger
+from .server import GraspServer
+from .session import GraspSession
+
+# Global server registry
+_servers: Dict[str, GraspServer] = {}
+
+async def _graceful_shutdown(reason: str = 'SIGTERM') -> None:
+    """Gracefully shutdown all servers.
+    
+    Args:
+        reason: Reason for shutdown
+    """
+    logger = get_logger().child('shutdown')
+    logger.info(f'Graceful shutdown initiated: {reason}')
+    
+    if not _servers:
+        logger.info('No active servers to shutdown')
+        return
+    
+    logger.info(f'Shutting down {len(_servers)} active servers...')
+    
+    # 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')
+
+
+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.SIGTERM, signal_handler)
+    signal.signal(signal.SIGINT, signal_handler)
+
+
+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()

grasp_sdk/models/__init__.py

@@ -23,6 +23,7 @@ class ISandboxConfig(TypedDict):
     timeout: int  # Required: Default timeout in milliseconds
     workspace: NotRequired[str]  # Optional: Grasp workspace ID
     debug: NotRequired[bool]  # Optional: Enable debug mode for detailed logging
+    keepAliveMs: NotRequired[int]  # Optional: Keep-alive time in milliseconds
     
 
 class IBrowserConfig(TypedDict):
@@ -35,19 +36,23 @@ class IBrowserConfig(TypedDict):
 
 class ICommandOptions(TypedDict):
     """Command execution options interface."""
-    inBackground: NotRequired[bool]  # Whether to run command in background
-    timeout: NotRequired[int]  # Timeout in milliseconds
+    inBackground: NotRequired[bool]  # Whether to run command in background (corresponds to 'background' in TypeScript)
+    timeout_ms: NotRequired[int]  # Timeout in milliseconds (corresponds to 'timeoutMs' in TypeScript)
     cwd: NotRequired[str]  # Working directory
-    nohup: NotRequired[bool]  # Whether to use nohup for command execution
+    user: NotRequired[str]  # User to run the command as (default: 'user')
+    envs: NotRequired[Dict[str, str]]  # Environment variables for the command
+    nohup: NotRequired[bool]  # Whether to use nohup for command execution (Grasp-specific extension)
+    # Note: onStdout and onStderr callbacks are handled differently in Python implementation
     
 
 class IScriptOptions(TypedDict):
     """Script execution options interface."""
-    type: str  # Required: Script type: 'cjs' for CommonJS, 'esm' for ES Modules
+    type: str  # Required: Script type: 'cjs' for CommonJS, 'esm' for ES Modules, 'py' for Python
     cwd: NotRequired[str]  # Working directory
-    timeoutMs: NotRequired[int]  # Timeout in milliseconds
+    timeout_ms: NotRequired[int]  # Timeout in milliseconds
     background: NotRequired[bool]  # Run in background
-    nohup: NotRequired[bool]  # Use nohup for background execution
+    user: NotRequired[str]  # User to run the script as (default: 'user')
+    nohup: NotRequired[bool]  # Use nohup for background execution (Grasp-specific extension)
     envs: NotRequired[Dict[str, str]]  # The environment variables
     preCommand: NotRequired[str]  # Pre command
 

grasp_sdk/services/__init__.py

@@ -4,5 +4,12 @@ This module contains core services for sandbox and browser management."""
 
 from .sandbox import SandboxService, CommandEventEmitter
 from .browser import BrowserService, CDPConnection
+from .terminal import TerminalService, Commands, StreamableCommandResult
+from .filesystem import FileSystemService
 
-__all__ = ['SandboxService', 'CommandEventEmitter', 'BrowserService', 'CDPConnection']
+__all__ = [
+    'SandboxService', 'CommandEventEmitter', 
+    'BrowserService', 'CDPConnection',
+    'TerminalService', 'Commands', 'StreamableCommandResult',
+    'FileSystemService'
+]