revert

Author: filip-michalsky

stagehand/main.py

@@ -391,9 +391,9 @@ async def init(self):
         self.logger.debug("Initializing Stagehand...")
         self.logger.debug(f"Environment: {self.env}")
 
-        # Always initialize playwright in a thread to avoid event loop conflicts
+        # Always initialize playwright with timeout to avoid hanging
         # This ensures compatibility with strict event loop environments like Langgraph
-        self._playwright = await self._init_playwright_in_thread()
+        self._playwright = await self._init_playwright_with_timeout()
 
         if self.env == "BROWSERBASE":
             # Create session if we don't have one
@@ -452,42 +452,38 @@ async def init(self):
 
         self._initialized = True
 
-    async def _init_playwright_in_thread(self):
+    async def _init_playwright_with_timeout(self):
         """
-        Initialize playwright using asyncio.to_thread to avoid blocking the main event loop.
+        Initialize playwright with a timeout to avoid hanging in strict event loop environments.
         
-        This method runs the potentially blocking async_playwright().start() in a thread
-        to avoid conflicts with strict event loop environments like Langgraph.
+        This method adds a timeout to the regular async_playwright().start() to prevent
+        hanging in environments like Langgraph that restrict blocking operations.
         """
-        self.logger.debug("Starting playwright initialization in background thread...")
-        
-        def _start_playwright_blocking():
-            """Start playwright in a blocking manner - will be run in a thread."""
-            import asyncio
-            from playwright.async_api import async_playwright
-            
-            # Create a new event loop for this thread
-            loop = asyncio.new_event_loop()
-            asyncio.set_event_loop(loop)
-            
-            try:
-                # Run the async playwright start in this loop
-                return loop.run_until_complete(async_playwright().start())
-            finally:
-                loop.close()
+        self.logger.debug("Starting playwright initialization with timeout...")
 
         try:
-            # Use asyncio.to_thread to run the blocking initialization
-            playwright_instance = await asyncio.to_thread(_start_playwright_blocking)
+            # Use asyncio.wait_for to add a timeout to prevent hanging
+            # If the environment doesn't allow blocking operations, this will fail fast
+            playwright_instance = await asyncio.wait_for(
+                async_playwright().start(),
+                timeout=30.0  # 30 second timeout
+            )
 
-            self.logger.debug("Playwright initialized successfully in background thread")
+            self.logger.debug("Playwright initialized successfully")
             return playwright_instance
 
+        except asyncio.TimeoutError:
+            self.logger.error("Playwright initialization timed out")
+            raise RuntimeError(
+                "Playwright initialization timed out after 30 seconds. This may indicate "
+                "your environment has strict event loop restrictions. If using Langgraph, "
+                "consider using the --allow-blocking flag."
+            )
         except Exception as e:
-            self.logger.error(f"Failed to initialize playwright in background thread: {e}")
+            self.logger.error(f"Failed to initialize playwright: {e}")
             raise RuntimeError(
-                "Failed to initialize Playwright in background thread. This may indicate a "
-                "deeper compatibility issue with your environment."
+                "Failed to initialize Playwright. This may indicate your environment "
+                "has restrictions on subprocess creation or event loop operations."
             ) from e
 
     def agent(self, **kwargs) -> Agent:

tests/unit/test_client_initialization.py

@@ -205,61 +205,79 @@ async def mock_create_session():
             await client._create_session()
 
     @pytest.mark.asyncio
-    @mock.patch("asyncio.to_thread")
-    async def test_init_playwright_in_thread(self, mock_to_thread):
-        """Test that playwright initialization works properly using asyncio.to_thread."""
+    @mock.patch("stagehand.main.async_playwright")
+    async def test_init_playwright_with_timeout(self, mock_async_playwright):
+        """Test that playwright initialization works properly with timeout."""
         # Create a mock playwright instance
         mock_playwright_instance = mock.AsyncMock()
         mock_playwright_instance.stop = mock.AsyncMock()
         mock_playwright_instance.chromium = mock.MagicMock()
         mock_playwright_instance.firefox = mock.MagicMock()
         mock_playwright_instance.webkit = mock.MagicMock()
 
-        # Mock asyncio.to_thread to return our mock instance
-        mock_to_thread.return_value = mock_playwright_instance
+        # Mock async_playwright().start() to return our mock instance as an awaitable
+        async def mock_start():
+            return mock_playwright_instance
+        
+        mock_async_playwright.return_value.start = mock_start
 
         # Create a Stagehand client with LOCAL env
         config = StagehandConfig(env="LOCAL")
         client = Stagehand(config=config)
 
-        # Test the threaded playwright initialization
-        result = await client._init_playwright_in_thread()
+        # Test the playwright initialization with timeout
+        result = await client._init_playwright_with_timeout()
 
         # Verify that the playwright instance was returned
         assert result is mock_playwright_instance
 
-        # Verify that asyncio.to_thread was called
-        mock_to_thread.assert_called_once()
-        
         # Verify the result has the expected attributes
         assert hasattr(result, 'chromium')
         assert hasattr(result, 'firefox') 
         assert hasattr(result, 'webkit')
         assert hasattr(result, 'stop')
 
     @pytest.mark.asyncio 
-    @mock.patch("asyncio.to_thread")
-    async def test_init_playwright_in_thread_handles_exceptions(self, mock_to_thread):
-        """Test that threaded playwright initialization properly handles exceptions."""
-        # Mock asyncio.to_thread to raise an exception
-        mock_to_thread.side_effect = Exception("Test exception")
+    @mock.patch("stagehand.main.async_playwright")
+    async def test_init_playwright_with_timeout_handles_exceptions(self, mock_async_playwright):
+        """Test that playwright initialization properly handles exceptions."""
+        # Mock async_playwright().start() to raise an exception as an awaitable
+        async def mock_start():
+            raise Exception("Test exception")
+        
+        mock_async_playwright.return_value.start = mock_start
 
         # Create a Stagehand client with LOCAL env
         config = StagehandConfig(env="LOCAL")
         client = Stagehand(config=config)
 
         # Test that the method raises a RuntimeError with our exception message
-        with pytest.raises(RuntimeError, match="Failed to initialize Playwright in background thread"):
-            await client._init_playwright_in_thread()
+        with pytest.raises(RuntimeError, match="Failed to initialize Playwright"):
+            await client._init_playwright_with_timeout()
+
+    @pytest.mark.asyncio 
+    @mock.patch("stagehand.main.asyncio.wait_for")
+    async def test_init_playwright_with_timeout_handles_timeout(self, mock_wait_for):
+        """Test that playwright initialization properly handles timeouts."""
+        # Mock asyncio.wait_for to raise a TimeoutError
+        mock_wait_for.side_effect = asyncio.TimeoutError()
+        
+        # Create a Stagehand client with LOCAL env
+        config = StagehandConfig(env="LOCAL")
+        client = Stagehand(config=config)
+        
+        # Test that the method raises a RuntimeError with timeout message
+        with pytest.raises(RuntimeError, match="Playwright initialization timed out"):
+            await client._init_playwright_with_timeout()
 
     @pytest.mark.asyncio
     @mock.patch("stagehand.main.cleanup_browser_resources")
     @mock.patch("stagehand.main.connect_local_browser")
-    @mock.patch.object(Stagehand, "_init_playwright_in_thread")
-    async def test_init_uses_threaded_playwright(
+    @mock.patch.object(Stagehand, "_init_playwright_with_timeout")
+    async def test_init_uses_playwright_with_timeout(
         self, mock_init_playwright, mock_connect_local, mock_cleanup
     ):
-        """Test that the main init() method uses threaded playwright initialization."""
+        """Test that the main init() method uses playwright initialization with timeout."""
         # Set up mocks
         mock_playwright_instance = mock.AsyncMock()
         mock_init_playwright.return_value = mock_playwright_instance
@@ -286,7 +304,7 @@ async def test_init_uses_threaded_playwright(
         # Initialize the client
         await client.init()
 
-        # Verify that threaded playwright initialization was called
+        # Verify that playwright initialization with timeout was called
         mock_init_playwright.assert_called_once()
 
         # Verify that the client is properly initialized