diff --git a/.changeset/gorilla-of-strongest-novelty.md b/.changeset/gorilla-of-strongest-novelty.md new file mode 100644 index 00000000..8d227ac5 --- /dev/null +++ b/.changeset/gorilla-of-strongest-novelty.md @@ -0,0 +1,5 @@ +--- +"stagehand": patch +--- + +Multi-tab support diff --git a/stagehand/browser.py b/stagehand/browser.py index 27278ed6..d21560d5 100644 --- a/stagehand/browser.py +++ b/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 diff --git a/stagehand/context.py b/stagehand/context.py index 51c0a06c..bfb6f2f7 100644 --- a/stagehand/context.py +++ b/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) diff --git a/stagehand/handlers/act_handler_utils.py b/stagehand/handlers/act_handler_utils.py index 947cf433..113366ac 100644 --- a/stagehand/handlers/act_handler_utils.py +++ b/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) diff --git a/stagehand/handlers/cua_handler.py b/stagehand/handlers/cua_handler.py index d718b1ba..2708aa3d 100644 --- a/stagehand/handlers/cua_handler.py +++ b/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: diff --git a/stagehand/main.py b/stagehand/main.py index ccdb8af8..45ee8c2b 100644 --- a/stagehand/main.py +++ b/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" {repr(stagehand._page)}>" + else: + return " 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,7 +501,7 @@ async def init(self): self._browser, self._context, self.context, - self.page, + self._page, ) = await connect_browserbase_browser( self._playwright, self.session_id, @@ -415,7 +509,7 @@ async def init(self): self, self.logger, ) - self._playwright_page = self.page._page + self._playwright_page = self._page._page except Exception: await self.close() raise @@ -427,7 +521,7 @@ 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, @@ -435,7 +529,7 @@ async def init(self): 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 diff --git a/tests/conftest.py b/tests/conftest.py index 4f29261f..36767e1a 100644 --- a/tests/conftest.py +++ b/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() diff --git a/tests/unit/core/test_live_page_proxy.py b/tests/unit/core/test_live_page_proxy.py new file mode 100644 index 00000000..d6a054bd --- /dev/null +++ b/tests/unit/core/test_live_page_proxy.py @@ -0,0 +1,199 @@ +"""Test the LivePageProxy functionality""" + +import asyncio +import pytest +from unittest.mock import AsyncMock, MagicMock, patch +from stagehand.main import LivePageProxy, Stagehand +from stagehand.page import StagehandPage + + +@pytest.mark.asyncio +async def test_live_page_proxy_basic_delegation(mock_stagehand_config): + """Test that LivePageProxy delegates to the active page""" + # Create a Stagehand instance + stagehand = Stagehand(config=mock_stagehand_config) + + # Mock page + mock_page = MagicMock(spec=StagehandPage) + mock_page.url = "https://active.com" + mock_page.title = AsyncMock(return_value="Active Page") + + # Set up the page + stagehand._page = mock_page + stagehand._initialized = True + + # Get the proxy + proxy = stagehand.page + + # Test that it delegates to the page + assert proxy.url == "https://active.com" + title = await proxy.title() + assert title == "Active Page" + + +@pytest.mark.asyncio +async def test_live_page_proxy_no_page_fallback(mock_stagehand_config): + """Test that LivePageProxy raises error when no page is set""" + # Create a Stagehand instance + stagehand = Stagehand(config=mock_stagehand_config) + + # No page set + stagehand._page = None + stagehand._initialized = True + + # Get the proxy + proxy = stagehand.page + + # Accessing attributes should raise RuntimeError + with pytest.raises(RuntimeError, match="No active page available"): + _ = proxy.url + + +@pytest.mark.asyncio +async def test_live_page_proxy_page_stability(mock_stagehand_config): + """Test that LivePageProxy waits for page stability on async operations""" + # Create a Stagehand instance + stagehand = Stagehand(config=mock_stagehand_config) + + # Track lock acquisition + lock_acquired = False + lock_released = False + + class TestLock: + async def __aenter__(self): + nonlocal lock_acquired + lock_acquired = True + await asyncio.sleep(0.1) # Simulate some work + return self + + async def __aexit__(self, *args): + nonlocal lock_released + lock_released = True + + stagehand._page_switch_lock = TestLock() + + # Mock page with async method + mock_page = MagicMock(spec=StagehandPage) + mock_page.click = AsyncMock(return_value=None) + + # Set up the page + stagehand._page = mock_page + stagehand._initialized = True + + # Get the proxy + proxy = stagehand.page + + # Call an async method (should wait for stability) + await proxy.click("button") + + # Verify lock was acquired and released + assert lock_acquired + assert lock_released + mock_page.click.assert_called_once_with("button") + + +@pytest.mark.asyncio +async def test_live_page_proxy_navigation_no_stability_check(mock_stagehand_config): + """Test that navigation methods don't wait for page stability""" + # Create a Stagehand instance + stagehand = Stagehand(config=mock_stagehand_config) + + # Track lock acquisition (should not happen) + lock_acquired = False + + class TestLock: + async def __aenter__(self): + nonlocal lock_acquired + lock_acquired = True + return self + + async def __aexit__(self, *args): + pass + + stagehand._page_switch_lock = TestLock() + + # Mock page with navigation methods + mock_page = MagicMock(spec=StagehandPage) + mock_page.goto = AsyncMock(return_value=None) + mock_page.reload = AsyncMock(return_value=None) + mock_page.go_back = AsyncMock(return_value=None) + mock_page.go_forward = AsyncMock(return_value=None) + + # Set up the page + stagehand._page = mock_page + stagehand._initialized = True + + # Get the proxy + proxy = stagehand.page + + # Call navigation methods (should NOT wait for stability) + await proxy.goto("https://example.com") + await proxy.reload() + await proxy.go_back() + await proxy.go_forward() + + # Verify lock was NOT acquired + assert not lock_acquired + + # Verify methods were called + mock_page.goto.assert_called_once_with("https://example.com") + mock_page.reload.assert_called_once() + mock_page.go_back.assert_called_once() + mock_page.go_forward.assert_called_once() + + +@pytest.mark.asyncio +async def test_live_page_proxy_dynamic_page_switching(mock_stagehand_config): + """Test that LivePageProxy dynamically switches between pages""" + # Create a Stagehand instance + stagehand = Stagehand(config=mock_stagehand_config) + + # Mock pages + page1 = MagicMock(spec=StagehandPage) + page1.url = "https://page1.com" + + page2 = MagicMock(spec=StagehandPage) + page2.url = "https://page2.com" + + # Set up initial state + stagehand._page = page1 + stagehand._initialized = True + + # Get the proxy + proxy = stagehand.page + + # Initially points to page1 + assert proxy.url == "https://page1.com" + + # Switch page + stagehand._page = page2 + + # Now points to page2 without creating a new proxy + assert proxy.url == "https://page2.com" + + +def test_live_page_proxy_no_page_error(mock_stagehand_config): + """Test that LivePageProxy raises error when no page is available""" + # Create a Stagehand instance + stagehand = Stagehand(config=mock_stagehand_config) + + # No page set + stagehand._page = None + stagehand._initialized = True + + # Get the proxy + proxy = stagehand.page + + # Accessing attributes should raise RuntimeError + with pytest.raises(RuntimeError, match="No active page available"): + _ = proxy.url + + +def test_live_page_proxy_not_initialized(mock_stagehand_config): + """Test that page property returns None when not initialized""" + # Create a Stagehand instance + stagehand = Stagehand(config=mock_stagehand_config) + stagehand._initialized = False + + # Should return None + assert stagehand.page is None \ No newline at end of file