revert to threading

Author: filip-michalsky

stagehand/main.py

@@ -210,26 +210,47 @@ def cleanup_handler(sig, frame):
                 # Try to get the current event loop
                 try:
                     loop = asyncio.get_running_loop()
+                    # Schedule cleanup in the existing event loop
+                    # Use call_soon_threadsafe since signal handlers run in a different thread context
+                    def schedule_cleanup():
+                        task = asyncio.create_task(self._async_cleanup())
+                        # Shield the task to prevent it from being cancelled
+                        asyncio.shield(task)
+                        # We don't need to await here since we're in call_soon_threadsafe
+
+                    loop.call_soon_threadsafe(schedule_cleanup)
+                    
                 except RuntimeError:
-                    # No event loop running - create one to run cleanup
-                    print("No event loop running, creating one for cleanup...")
+                    # No event loop running - use thread-based cleanup to avoid blocking
+                    print("No event loop running, scheduling cleanup in separate thread...")
                     try:
-                        asyncio.run(self._async_cleanup())
+                        # Use threading to avoid blocking operations in signal handler
+                        import threading
+                        
+                        def cleanup_in_thread():
+                            try:
+                                # Create new event loop in thread for cleanup
+                                new_loop = asyncio.new_event_loop()
+                                asyncio.set_event_loop(new_loop)
+                                new_loop.run_until_complete(self._async_cleanup())
+                            except Exception as e:
+                                print(f"Error during threaded cleanup: {str(e)}")
+                            finally:
+                                try:
+                                    new_loop.close()
+                                except Exception:
+                                    pass
+                                sys.exit(0)
+                        
+                        cleanup_thread = threading.Thread(target=cleanup_in_thread, daemon=True)
+                        cleanup_thread.start()
+                        # Give the cleanup thread time to start, but don't wait indefinitely
+                        cleanup_thread.join(timeout=5.0)
+                        sys.exit(0)
+                        
                     except Exception as e:
                         print(f"Error during cleanup: {str(e)}")
-                    finally:
-                        sys.exit(0)
-                    return
-
-                # Schedule cleanup in the existing event loop
-                # Use call_soon_threadsafe since signal handlers run in a different thread context
-                def schedule_cleanup():
-                    task = asyncio.create_task(self._async_cleanup())
-                    # Shield the task to prevent it from being cancelled
-                    asyncio.shield(task)
-                    # We don't need to await here since we're in call_soon_threadsafe
-
-                loop.call_soon_threadsafe(schedule_cleanup)
+                        sys.exit(1)
 
             except Exception as e:
                 print(f"Error during signal cleanup: {str(e)}")
@@ -454,34 +475,40 @@ async def init(self):
 
     async def _init_playwright_with_timeout(self):
         """
-        Initialize playwright with a timeout to avoid hanging in strict event loop environments.
+        Initialize playwright in a separate thread to avoid blocking operations in strict event loop environments.
 
-        This method adds a timeout to the regular async_playwright().start() to prevent
-        hanging in environments that restrict blocking operations.
+        This method uses asyncio.to_thread to run Playwright initialization in a separate thread
+        with its own event loop, preventing blocking operations from affecting the main event loop
+        in environments like Langgraph that enforce strict async policies.
         """
-        self.logger.debug("Starting playwright initialization with timeout...")
+        self.logger.debug("Starting playwright initialization in separate thread...")
 
         try:
-            # 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
-            )
+            # Run Playwright initialization in a separate thread to avoid blocking the main event loop
+            # This is necessary for environments like Langgraph that have strict blocking operation restrictions
+            def _init_playwright_in_thread():
+                # Create a new event loop for this thread
+                import asyncio
+                new_loop = asyncio.new_event_loop()
+                asyncio.set_event_loop(new_loop)
+                
+                try:
+                    # Run the Playwright initialization in this thread's event loop
+                    return new_loop.run_until_complete(async_playwright().start())
+                finally:
+                    new_loop.close()
 
-            self.logger.debug("Playwright initialized successfully")
+            # Use asyncio.to_thread to run the blocking initialization in a separate thread
+            playwright_instance = await asyncio.to_thread(_init_playwright_in_thread)
+
+            self.logger.debug("Playwright initialized successfully in separate thread")
             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."
-            ) from None
         except Exception as e:
-            self.logger.error(f"Failed to initialize playwright: {e}")
+            self.logger.error(f"Failed to initialize playwright in thread: {e}")
             raise RuntimeError(
                 "Failed to initialize Playwright. This may indicate your environment "
-                "has restrictions on subprocess creation or event loop operations."
+                "has restrictions on subprocess creation or threading operations."
             ) from e
 
     def agent(self, **kwargs) -> Agent: