Skip to content

Commit f68e86c

Browse files
miguelg719arunpatroArun Patrofilip-michalsky
authored
feat: don't close new opened tabs (#161) (#169)
* feat: don't close new opened tabs (#161) * feat: port new page handling from JS Stagehand PR #844 - Add live page proxy that dynamically tracks the focused page - Implement context event listener to initialize new pages automatically - Remove automatic tab closing behavior in act_handler_utils and cua_handler - Keep both original and new tabs open when new pages are created - Ensure stagehand.page always references the current active page This implementation matches the behavior of browserbase/stagehand#844 * Update stagehand/handlers/cua_handler.py * Update stagehand/handlers/act_handler_utils.py --------- Co-authored-by: Arun Patro <[email protected]> Co-authored-by: Miguel <[email protected]> * formatting, logs * changeset * fix: update mock_stagehand_client fixture to set internal page properties The page property is now read-only (returns LivePageProxy), so tests need to set the internal _original_page and _active_page properties instead * feat: add page stability check to LivePageProxy for async operations Ensures async operations wait for any pending page switches to complete * formatting * fix: prevent deadlock in page navigation and add page stability tests - Initialize _page_switch_lock in Stagehand constructor - Skip page stability check for navigation methods (goto, reload, go_back, go_forward) - Use lock when switching active pages in context - Add comprehensive tests for LivePageProxy functionality * consolidate original and active page to just one page * Update stagehand/context.py * Update .changeset/gorilla-of-strongest-novelty.md * timeouts * python 3.10 or less compatibility --------- Co-authored-by: Arun Patro <[email protected]> Co-authored-by: Arun Patro <[email protected]> Co-authored-by: Filip Michalsky <[email protected]>
1 parent 021c946 commit f68e86c

File tree

8 files changed

+394
-24
lines changed

8 files changed

+394
-24
lines changed
Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
---
2+
"stagehand": patch
3+
---
4+
5+
Multi-tab support

stagehand/browser.py

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -242,11 +242,10 @@ async def connect_local_browser(
242242
if context.pages:
243243
playwright_page = context.pages[0]
244244
logger.debug("Using initial page from local context.")
245+
page = await stagehand_context.get_stagehand_page(playwright_page)
245246
else:
246247
logger.debug("No initial page found, creating a new one.")
247-
playwright_page = await context.new_page()
248-
249-
page = StagehandPage(playwright_page, stagehand_instance)
248+
page = await stagehand_context.new_page()
250249

251250
return browser, context, stagehand_context, page, temp_user_data_dir
252251

stagehand/context.py

Lines changed: 53 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import asyncio
12
import os
23
import weakref
34

@@ -40,7 +41,8 @@ async def inject_custom_scripts(self, pw_page: Page):
4041
async def get_stagehand_page(self, pw_page: Page) -> StagehandPage:
4142
if pw_page not in self.page_map:
4243
return await self.create_stagehand_page(pw_page)
43-
return self.page_map[pw_page]
44+
stagehand_page = self.page_map[pw_page]
45+
return stagehand_page
4446

4547
async def get_stagehand_pages(self) -> list:
4648
# Return a list of StagehandPage wrappers for all pages in the context
@@ -53,25 +55,74 @@ async def get_stagehand_pages(self) -> list:
5355

5456
def set_active_page(self, stagehand_page: StagehandPage):
5557
self.active_stagehand_page = stagehand_page
56-
# Optionally update the active page in the stagehand client if needed
58+
# Update the active page in the stagehand client
5759
if hasattr(self.stagehand, "_set_active_page"):
5860
self.stagehand._set_active_page(stagehand_page)
61+
self.stagehand.logger.debug(
62+
f"Set active page to: {stagehand_page.url}", category="context"
63+
)
64+
else:
65+
self.stagehand.logger.debug(
66+
"Stagehand does not have _set_active_page method", category="context"
67+
)
5968

6069
def get_active_page(self) -> StagehandPage:
6170
return self.active_stagehand_page
6271

6372
@classmethod
6473
async def init(cls, context: BrowserContext, stagehand):
74+
stagehand.logger.debug("StagehandContext.init() called", category="context")
6575
instance = cls(context, stagehand)
6676
# Pre-initialize StagehandPages for any existing pages
77+
stagehand.logger.debug(
78+
f"Found {len(instance._context.pages)} existing pages", category="context"
79+
)
6780
for pw_page in instance._context.pages:
6881
await instance.create_stagehand_page(pw_page)
6982
if instance._context.pages:
7083
first_page = instance._context.pages[0]
7184
stagehand_page = await instance.get_stagehand_page(first_page)
7285
instance.set_active_page(stagehand_page)
86+
87+
# Add event listener for new pages (popups, new tabs from window.open, etc.)
88+
def handle_page_event(pw_page):
89+
# Playwright expects sync handler, so we schedule the async work
90+
asyncio.create_task(instance._handle_new_page(pw_page))
91+
92+
context.on("page", handle_page_event)
93+
7394
return instance
7495

96+
async def _handle_new_page(self, pw_page: Page):
97+
"""
98+
Handle new pages created by the browser (popups, window.open, etc.).
99+
Uses the page switch lock to prevent race conditions with ongoing operations.
100+
"""
101+
try:
102+
# Use wait_for for Python 3.10 compatibility (timeout prevents indefinite blocking)
103+
async def handle_with_lock():
104+
async with self.stagehand._page_switch_lock:
105+
self.stagehand.logger.debug(
106+
f"Creating StagehandPage for new page with URL: {pw_page.url}",
107+
category="context",
108+
)
109+
stagehand_page = await self.create_stagehand_page(pw_page)
110+
self.set_active_page(stagehand_page)
111+
self.stagehand.logger.debug(
112+
"New page detected and initialized", category="context"
113+
)
114+
115+
await asyncio.wait_for(handle_with_lock(), timeout=30)
116+
except asyncio.TimeoutError:
117+
self.stagehand.logger.error(
118+
f"Timeout waiting for page switch lock when handling new page: {pw_page.url}",
119+
category="context",
120+
)
121+
except Exception as e:
122+
self.stagehand.logger.error(
123+
f"Failed to initialize new page: {str(e)}", category="context"
124+
)
125+
75126
def __getattr__(self, name):
76127
# Forward attribute lookups to the underlying BrowserContext
77128
attr = getattr(self._context, name)

stagehand/handlers/act_handler_utils.py

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -471,10 +471,6 @@ async def handle_possible_page_navigation(
471471
category="action",
472472
auxiliary={"url": {"value": new_opened_tab.url, "type": "string"}},
473473
)
474-
new_tab_url = new_opened_tab.url
475-
await new_opened_tab.close()
476-
await stagehand_page._page.goto(new_tab_url)
477-
await stagehand_page._page.wait_for_load_state("domcontentloaded")
478474

479475
try:
480476
await stagehand_page._wait_for_settled_dom(dom_settle_timeout_ms)

stagehand/handlers/cua_handler.py

Lines changed: 7 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -559,21 +559,19 @@ async def handle_page_navigation(
559559
pass # The action that might open a page has already run. We check if one was caught.
560560
newly_opened_page = await new_page_info.value
561561

562-
new_page_url = newly_opened_page.url
563-
await newly_opened_page.close()
564-
await self.page.goto(new_page_url, timeout=dom_settle_timeout_ms)
565-
# After navigating, the DOM needs to settle on the new URL.
566-
await self._wait_for_settled_dom(timeout_ms=dom_settle_timeout_ms)
562+
# Don't close the new tab - let it remain open and be handled by the context
563+
# The StagehandContext will automatically make this the active page via its event listener
564+
self.logger.debug(
565+
f"New page detected with URL: {newly_opened_page.url}",
566+
category=StagehandFunctionName.AGENT,
567+
)
567568

568569
except asyncio.TimeoutError:
569570
newly_opened_page = None
570571
except Exception:
571572
newly_opened_page = None
572573

573-
# If no new tab was opened and handled by navigating, or if we are on the original page after handling a new tab,
574-
# then proceed to wait for DOM settlement on the current page.
575-
if not newly_opened_page:
576-
await self._wait_for_settled_dom(timeout_ms=dom_settle_timeout_ms)
574+
await self._wait_for_settled_dom(timeout_ms=dom_settle_timeout_ms)
577575

578576
final_url = self.page.url
579577
if final_url != initial_url:

stagehand/main.py

Lines changed: 126 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,98 @@
3333
load_dotenv()
3434

3535

36+
class LivePageProxy:
37+
"""
38+
A proxy object that dynamically delegates all operations to the current active page.
39+
This mimics the behavior of the JavaScript Proxy in the original implementation.
40+
"""
41+
42+
def __init__(self, stagehand_instance):
43+
# Use object.__setattr__ to avoid infinite recursion
44+
object.__setattr__(self, "_stagehand", stagehand_instance)
45+
46+
async def _ensure_page_stability(self):
47+
"""Wait for any pending page switches to complete"""
48+
if hasattr(self._stagehand, "_page_switch_lock"):
49+
try:
50+
# Use wait_for for Python 3.10 compatibility (timeout prevents indefinite blocking)
51+
async def acquire_lock():
52+
async with self._stagehand._page_switch_lock:
53+
pass # Just wait for any ongoing switches
54+
55+
await asyncio.wait_for(acquire_lock(), timeout=30)
56+
except asyncio.TimeoutError:
57+
# Log the timeout and raise to let caller handle it
58+
if hasattr(self._stagehand, "logger"):
59+
self._stagehand.logger.error(
60+
"Timeout waiting for page stability lock", category="live_proxy"
61+
)
62+
raise RuntimeError from asyncio.TimeoutError(
63+
"Page stability lock timeout - possible deadlock detected"
64+
)
65+
66+
def __getattr__(self, name):
67+
"""Delegate all attribute access to the current active page."""
68+
stagehand = object.__getattribute__(self, "_stagehand")
69+
70+
# Get the current page
71+
if hasattr(stagehand, "_page") and stagehand._page:
72+
page = stagehand._page
73+
else:
74+
raise RuntimeError("No active page available")
75+
76+
# For async operations, make them wait for stability
77+
attr = getattr(page, name)
78+
if callable(attr) and asyncio.iscoroutinefunction(attr):
79+
# Don't wait for stability on navigation methods
80+
if name in ["goto", "reload", "go_back", "go_forward"]:
81+
return attr
82+
83+
async def wrapped(*args, **kwargs):
84+
await self._ensure_page_stability()
85+
return await attr(*args, **kwargs)
86+
87+
return wrapped
88+
return attr
89+
90+
def __setattr__(self, name, value):
91+
"""Delegate all attribute setting to the current active page."""
92+
if name.startswith("_"):
93+
# Internal attributes are set on the proxy itself
94+
object.__setattr__(self, name, value)
95+
else:
96+
stagehand = object.__getattribute__(self, "_stagehand")
97+
98+
# Get the current page
99+
if hasattr(stagehand, "_page") and stagehand._page:
100+
page = stagehand._page
101+
else:
102+
raise RuntimeError("No active page available")
103+
104+
# Set the attribute on the page
105+
setattr(page, name, value)
106+
107+
def __dir__(self):
108+
"""Return attributes of the current active page."""
109+
stagehand = object.__getattribute__(self, "_stagehand")
110+
111+
if hasattr(stagehand, "_page") and stagehand._page:
112+
page = stagehand._page
113+
else:
114+
return []
115+
116+
return dir(page)
117+
118+
def __repr__(self):
119+
"""Return representation of the current active page."""
120+
stagehand = object.__getattribute__(self, "_stagehand")
121+
122+
if hasattr(stagehand, "_page") and stagehand._page:
123+
return f"<LivePageProxy -> {repr(stagehand._page)}>"
124+
else:
125+
return "<LivePageProxy -> No active page>"
126+
127+
36128
class Stagehand:
37129
"""
38130
Main Stagehand class.
@@ -166,7 +258,7 @@ def __init__(
166258
self._browser = None
167259
self._context: Optional[BrowserContext] = None
168260
self._playwright_page: Optional[PlaywrightPage] = None
169-
self.page: Optional[StagehandPage] = None
261+
self._page: Optional[StagehandPage] = None
170262
self.context: Optional[StagehandContext] = None
171263
self.use_api = self.config.use_api
172264
self.experimental = self.config.experimental
@@ -181,6 +273,8 @@ def __init__(
181273

182274
self._initialized = False # Flag to track if init() has run
183275
self._closed = False # Flag to track if resources have been closed
276+
self._live_page_proxy = None # Live page proxy
277+
self._page_switch_lock = asyncio.Lock() # Lock for page stability
184278

185279
# Setup LLM client if LOCAL mode
186280
self.llm = None
@@ -407,15 +501,15 @@ async def init(self):
407501
self._browser,
408502
self._context,
409503
self.context,
410-
self.page,
504+
self._page,
411505
) = await connect_browserbase_browser(
412506
self._playwright,
413507
self.session_id,
414508
self.browserbase_api_key,
415509
self,
416510
self.logger,
417511
)
418-
self._playwright_page = self.page._page
512+
self._playwright_page = self._page._page
419513
except Exception:
420514
await self.close()
421515
raise
@@ -427,15 +521,15 @@ async def init(self):
427521
self._browser,
428522
self._context,
429523
self.context,
430-
self.page,
524+
self._page,
431525
self._local_user_data_dir_temp,
432526
) = await connect_local_browser(
433527
self._playwright,
434528
self.local_browser_launch_options,
435529
self,
436530
self.logger,
437531
)
438-
self._playwright_page = self.page._page
532+
self._playwright_page = self._page._page
439533
except Exception:
440534
await self.close()
441535
raise
@@ -615,6 +709,33 @@ def _handle_llm_metrics(
615709

616710
self.update_metrics_from_response(function_enum, response, inference_time_ms)
617711

712+
def _set_active_page(self, stagehand_page: StagehandPage):
713+
"""
714+
Internal method called by StagehandContext to update the active page.
715+
716+
Args:
717+
stagehand_page: The StagehandPage to set as active
718+
"""
719+
self._page = stagehand_page
720+
721+
@property
722+
def page(self) -> Optional[StagehandPage]:
723+
"""
724+
Get the current active page. This property returns a live proxy that
725+
always points to the currently focused page when multiple tabs are open.
726+
727+
Returns:
728+
A LivePageProxy that delegates to the active StagehandPage or None if not initialized
729+
"""
730+
if not self._initialized:
731+
return None
732+
733+
# Create the live page proxy if it doesn't exist
734+
if not self._live_page_proxy:
735+
self._live_page_proxy = LivePageProxy(self)
736+
737+
return self._live_page_proxy
738+
618739

619740
# Bind the imported API methods to the Stagehand class
620741
Stagehand._create_session = _create_session

tests/conftest.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -233,7 +233,8 @@ def mock_stagehand_client(mock_stagehand_config):
233233
# Mock the essential components
234234
client.llm = MagicMock()
235235
client.llm.completion = AsyncMock()
236-
client.page = MagicMock()
236+
# Set internal page property instead of the read-only page property
237+
client._page = MagicMock()
237238
client.agent = MagicMock()
238239
client._client = MagicMock()
239240
client._execute = AsyncMock()

0 commit comments

Comments
 (0)