feat: don't close new opened tabs (#161)

.changeset/gorilla-of-strongest-novelty.md

@@ -0,0 +1,5 @@
+---
+"stagehand": patch
+---
+
+Multi-tab support

stagehand/browser.py

@@ -242,11 +242,10 @@ async def connect_local_browser(
     if context.pages:
         playwright_page = context.pages[0]
         logger.debug("Using initial page from local context.")
+        page = await stagehand_context.get_stagehand_page(playwright_page)
     else:
         logger.debug("No initial page found, creating a new one.")
-        playwright_page = await context.new_page()
-
-    page = StagehandPage(playwright_page, stagehand_instance)
+        page = await stagehand_context.new_page()
 
     return browser, context, stagehand_context, page, temp_user_data_dir
 

stagehand/context.py

@@ -1,3 +1,4 @@
+import asyncio
 import os
 import weakref
 
@@ -40,7 +41,8 @@ async def inject_custom_scripts(self, pw_page: Page):
     async def get_stagehand_page(self, pw_page: Page) -> StagehandPage:
         if pw_page not in self.page_map:
             return await self.create_stagehand_page(pw_page)
-        return self.page_map[pw_page]
+        stagehand_page = self.page_map[pw_page]
+        return stagehand_page
 
     async def get_stagehand_pages(self) -> list:
         # Return a list of StagehandPage wrappers for all pages in the context
@@ -53,25 +55,74 @@ async def get_stagehand_pages(self) -> list:
 
     def set_active_page(self, stagehand_page: StagehandPage):
         self.active_stagehand_page = stagehand_page
-        # Optionally update the active page in the stagehand client if needed
+        # Update the active page in the stagehand client
         if hasattr(self.stagehand, "_set_active_page"):
             self.stagehand._set_active_page(stagehand_page)
+            self.stagehand.logger.debug(
+                f"Set active page to: {stagehand_page.url}", category="context"
+            )
+        else:
+            self.stagehand.logger.debug(
+                "Stagehand does not have _set_active_page method", category="context"
+            )
 
     def get_active_page(self) -> StagehandPage:
         return self.active_stagehand_page
 
     @classmethod
     async def init(cls, context: BrowserContext, stagehand):
+        stagehand.logger.debug("StagehandContext.init() called", category="context")
         instance = cls(context, stagehand)
         # Pre-initialize StagehandPages for any existing pages
+        stagehand.logger.debug(
+            f"Found {len(instance._context.pages)} existing pages", category="context"
+        )
         for pw_page in instance._context.pages:
             await instance.create_stagehand_page(pw_page)
         if instance._context.pages:
             first_page = instance._context.pages[0]
             stagehand_page = await instance.get_stagehand_page(first_page)
             instance.set_active_page(stagehand_page)
+
+        # Add event listener for new pages (popups, new tabs from window.open, etc.)
+        def handle_page_event(pw_page):
+            # Playwright expects sync handler, so we schedule the async work
+            asyncio.create_task(instance._handle_new_page(pw_page))
+
+        context.on("page", handle_page_event)
+
         return instance
 
+    async def _handle_new_page(self, pw_page: Page):
+        """
+        Handle new pages created by the browser (popups, window.open, etc.).
+        Uses the page switch lock to prevent race conditions with ongoing operations.
+        """
+        try:
+            # Use wait_for for Python 3.10 compatibility (timeout prevents indefinite blocking)
+            async def handle_with_lock():
+                async with self.stagehand._page_switch_lock:
+                    self.stagehand.logger.debug(
+                        f"Creating StagehandPage for new page with URL: {pw_page.url}",
+                        category="context",
+                    )
+                    stagehand_page = await self.create_stagehand_page(pw_page)
+                    self.set_active_page(stagehand_page)
+                    self.stagehand.logger.debug(
+                        "New page detected and initialized", category="context"
+                    )
+
+            await asyncio.wait_for(handle_with_lock(), timeout=30)
+        except asyncio.TimeoutError:
+            self.stagehand.logger.error(
+                f"Timeout waiting for page switch lock when handling new page: {pw_page.url}",
+                category="context",
+            )
+        except Exception as e:
+            self.stagehand.logger.error(
+                f"Failed to initialize new page: {str(e)}", category="context"
+            )
+
     def __getattr__(self, name):
         # Forward attribute lookups to the underlying BrowserContext
         attr = getattr(self._context, name)

stagehand/handlers/act_handler_utils.py

@@ -471,10 +471,6 @@ async def handle_possible_page_navigation(
             category="action",
             auxiliary={"url": {"value": new_opened_tab.url, "type": "string"}},
         )
-        new_tab_url = new_opened_tab.url
-        await new_opened_tab.close()
-        await stagehand_page._page.goto(new_tab_url)
-        await stagehand_page._page.wait_for_load_state("domcontentloaded")
 
     try:
         await stagehand_page._wait_for_settled_dom(dom_settle_timeout_ms)

stagehand/handlers/cua_handler.py

@@ -559,21 +559,19 @@ async def handle_page_navigation(
                 pass  # The action that might open a page has already run. We check if one was caught.
             newly_opened_page = await new_page_info.value
 
-            new_page_url = newly_opened_page.url
-            await newly_opened_page.close()
-            await self.page.goto(new_page_url, timeout=dom_settle_timeout_ms)
-            # After navigating, the DOM needs to settle on the new URL.
-            await self._wait_for_settled_dom(timeout_ms=dom_settle_timeout_ms)
+            # Don't close the new tab - let it remain open and be handled by the context
+            # The StagehandContext will automatically make this the active page via its event listener
+            self.logger.debug(
+                f"New page detected with URL: {newly_opened_page.url}",
+                category=StagehandFunctionName.AGENT,
+            )
 
         except asyncio.TimeoutError:
             newly_opened_page = None
         except Exception:
             newly_opened_page = None
 
-        # If no new tab was opened and handled by navigating, or if we are on the original page after handling a new tab,
-        # then proceed to wait for DOM settlement on the current page.
-        if not newly_opened_page:
-            await self._wait_for_settled_dom(timeout_ms=dom_settle_timeout_ms)
+        await self._wait_for_settled_dom(timeout_ms=dom_settle_timeout_ms)
 
         final_url = self.page.url
         if final_url != initial_url:

stagehand/main.py

@@ -33,6 +33,98 @@
 load_dotenv()
 
 
+class LivePageProxy:
+    """
+    A proxy object that dynamically delegates all operations to the current active page.
+    This mimics the behavior of the JavaScript Proxy in the original implementation.
+    """
+
+    def __init__(self, stagehand_instance):
+        # Use object.__setattr__ to avoid infinite recursion
+        object.__setattr__(self, "_stagehand", stagehand_instance)
+
+    async def _ensure_page_stability(self):
+        """Wait for any pending page switches to complete"""
+        if hasattr(self._stagehand, "_page_switch_lock"):
+            try:
+                # Use wait_for for Python 3.10 compatibility (timeout prevents indefinite blocking)
+                async def acquire_lock():
+                    async with self._stagehand._page_switch_lock:
+                        pass  # Just wait for any ongoing switches
+
+                await asyncio.wait_for(acquire_lock(), timeout=30)
+            except asyncio.TimeoutError:
+                # Log the timeout and raise to let caller handle it
+                if hasattr(self._stagehand, "logger"):
+                    self._stagehand.logger.error(
+                        "Timeout waiting for page stability lock", category="live_proxy"
+                    )
+                raise RuntimeError from asyncio.TimeoutError(
+                    "Page stability lock timeout - possible deadlock detected"
+                )
+
+    def __getattr__(self, name):
+        """Delegate all attribute access to the current active page."""
+        stagehand = object.__getattribute__(self, "_stagehand")
+
+        # Get the current page
+        if hasattr(stagehand, "_page") and stagehand._page:
+            page = stagehand._page
+        else:
+            raise RuntimeError("No active page available")
+
+        # For async operations, make them wait for stability
+        attr = getattr(page, name)
+        if callable(attr) and asyncio.iscoroutinefunction(attr):
+            # Don't wait for stability on navigation methods
+            if name in ["goto", "reload", "go_back", "go_forward"]:
+                return attr
+
+            async def wrapped(*args, **kwargs):
+                await self._ensure_page_stability()
+                return await attr(*args, **kwargs)
+
+            return wrapped
+        return attr
+
+    def __setattr__(self, name, value):
+        """Delegate all attribute setting to the current active page."""
+        if name.startswith("_"):
+            # Internal attributes are set on the proxy itself
+            object.__setattr__(self, name, value)
+        else:
+            stagehand = object.__getattribute__(self, "_stagehand")
+
+            # Get the current page
+            if hasattr(stagehand, "_page") and stagehand._page:
+                page = stagehand._page
+            else:
+                raise RuntimeError("No active page available")
+
+            # Set the attribute on the page
+            setattr(page, name, value)
+
+    def __dir__(self):
+        """Return attributes of the current active page."""
+        stagehand = object.__getattribute__(self, "_stagehand")
+
+        if hasattr(stagehand, "_page") and stagehand._page:
+            page = stagehand._page
+        else:
+            return []
+
+        return dir(page)
+
+    def __repr__(self):
+        """Return representation of the current active page."""
+        stagehand = object.__getattribute__(self, "_stagehand")
+
+        if hasattr(stagehand, "_page") and stagehand._page:
+            return f"<LivePageProxy -> {repr(stagehand._page)}>"
+        else:
+            return "<LivePageProxy -> No active page>"
+
+
 class Stagehand:
     """
     Main Stagehand class.
@@ -166,7 +258,7 @@ def __init__(
         self._browser = None
         self._context: Optional[BrowserContext] = None
         self._playwright_page: Optional[PlaywrightPage] = None
-        self.page: Optional[StagehandPage] = None
+        self._page: Optional[StagehandPage] = None
         self.context: Optional[StagehandContext] = None
         self.use_api = self.config.use_api
         self.experimental = self.config.experimental
@@ -181,6 +273,8 @@ def __init__(
 
         self._initialized = False  # Flag to track if init() has run
         self._closed = False  # Flag to track if resources have been closed
+        self._live_page_proxy = None  # Live page proxy
+        self._page_switch_lock = asyncio.Lock()  # Lock for page stability
 
         # Setup LLM client if LOCAL mode
         self.llm = None
@@ -407,15 +501,15 @@ async def init(self):
                     self._browser,
                     self._context,
                     self.context,
-                    self.page,
+                    self._page,
                 ) = await connect_browserbase_browser(
                     self._playwright,
                     self.session_id,
                     self.browserbase_api_key,
                     self,
                     self.logger,
                 )
-                self._playwright_page = self.page._page
+                self._playwright_page = self._page._page
             except Exception:
                 await self.close()
                 raise
@@ -427,15 +521,15 @@ async def init(self):
                     self._browser,
                     self._context,
                     self.context,
-                    self.page,
+                    self._page,
                     self._local_user_data_dir_temp,
                 ) = await connect_local_browser(
                     self._playwright,
                     self.local_browser_launch_options,
                     self,
                     self.logger,
                 )
-                self._playwright_page = self.page._page
+                self._playwright_page = self._page._page
             except Exception:
                 await self.close()
                 raise
@@ -615,6 +709,33 @@ def _handle_llm_metrics(
 
         self.update_metrics_from_response(function_enum, response, inference_time_ms)
 
+    def _set_active_page(self, stagehand_page: StagehandPage):
+        """
+        Internal method called by StagehandContext to update the active page.
+
+        Args:
+            stagehand_page: The StagehandPage to set as active
+        """
+        self._page = stagehand_page
+
+    @property
+    def page(self) -> Optional[StagehandPage]:
+        """
+        Get the current active page. This property returns a live proxy that
+        always points to the currently focused page when multiple tabs are open.
+
+        Returns:
+            A LivePageProxy that delegates to the active StagehandPage or None if not initialized
+        """
+        if not self._initialized:
+            return None
+
+        # Create the live page proxy if it doesn't exist
+        if not self._live_page_proxy:
+            self._live_page_proxy = LivePageProxy(self)
+
+        return self._live_page_proxy
+
 
 # Bind the imported API methods to the Stagehand class
 Stagehand._create_session = _create_session

tests/conftest.py

@@ -233,7 +233,8 @@ def mock_stagehand_client(mock_stagehand_config):
         # Mock the essential components
         client.llm = MagicMock()
         client.llm.completion = AsyncMock()
-        client.page = MagicMock()
+        # Set internal page property instead of the read-only page property
+        client._page = MagicMock()
         client.agent = MagicMock()
         client._client = MagicMock()
         client._execute = AsyncMock()