event loop is blocked - fixing

Author: filip-michalsky

stagehand/main.py

@@ -454,38 +454,39 @@ async def init(self):
 
     async def _init_playwright_in_thread(self):
         """
-        Initialize playwright in a separate thread for compatibility with strict event loops.
+        Initialize playwright using asyncio.to_thread to avoid blocking the main event loop.
         
-        This method runs the playwright initialization in a separate thread to avoid
-        blocking operations that can conflict with strict event loop environments
-        like Langgraph when running without --allow-blocking.
+        This method runs the potentially blocking async_playwright().start() in a thread
+        to avoid conflicts with strict event loop environments like Langgraph.
         """
-        def _start_playwright():
-            """Helper function to start playwright in a separate thread."""
+        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:
-                # Start playwright in this thread's event loop
+                # Run the async playwright start in this loop
                 return loop.run_until_complete(async_playwright().start())
             finally:
-                # Clean up the loop
                 loop.close()
 
-        self.logger.debug("Starting playwright in separate thread...")
-        
         try:
-            # Run the playwright initialization in a thread
-            playwright_instance = await asyncio.to_thread(_start_playwright)
-            self.logger.debug("Playwright initialized successfully in thread")
+            # Use asyncio.to_thread to run the blocking initialization
+            playwright_instance = await asyncio.to_thread(_start_playwright_blocking)
+            
+            self.logger.debug("Playwright initialized successfully in background thread")
             return playwright_instance
+            
         except Exception as e:
-            self.logger.error(f"Failed to initialize playwright in thread: {e}")
+            self.logger.error(f"Failed to initialize playwright in background thread: {e}")
             raise RuntimeError(
-                "Failed to initialize Playwright in thread. This may indicate a "
+                "Failed to initialize Playwright in background thread. This may indicate a "
                 "deeper compatibility issue with your environment."
             ) from e
 

tests/unit/test_client_initialization.py

@@ -205,17 +205,18 @@ async def mock_create_session():
             await client._create_session()
 
     @pytest.mark.asyncio
-    @mock.patch("stagehand.main.async_playwright")
-    async def test_init_playwright_in_thread(self, mock_async_playwright):
-        """Test that playwright initialization works properly in a separate thread."""
+    @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."""
         # 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 the async_playwright().start() to return our mock instance
-        mock_async_playwright_start = mock.AsyncMock(return_value=mock_playwright_instance)
-        mock_async_playwright.return_value.start = mock_async_playwright_start
+        # Mock asyncio.to_thread to return our mock instance
+        mock_to_thread.return_value = mock_playwright_instance
 
         # Create a Stagehand client with LOCAL env
         config = StagehandConfig(env="LOCAL")
@@ -227,26 +228,28 @@ async def test_init_playwright_in_thread(self, mock_async_playwright):
         # Verify that the playwright instance was returned
         assert result is mock_playwright_instance
 
-        # Verify that async_playwright().start() was called
-        mock_async_playwright_start.assert_called_once()
+        # 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("stagehand.main.async_playwright")
-    async def test_init_playwright_in_thread_handles_exceptions(self, mock_async_playwright):
+    @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 async_playwright().start() to raise an exception
-        mock_async_playwright.return_value.start.side_effect = Exception("Test exception")
+        # Mock asyncio.to_thread to raise an exception
+        mock_to_thread.side_effect = Exception("Test exception")
 
         # 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 thread"):
+        with pytest.raises(RuntimeError, match="Failed to initialize Playwright in background thread"):
             await client._init_playwright_in_thread()
 
     @pytest.mark.asyncio