-
Notifications
You must be signed in to change notification settings - Fork 61
feat: don't close new opened tabs (#161) #169
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from 10 commits
78525a7
b2ee152
535339d
a474768
bbd8289
55800e4
29e34e2
c0959a2
c51f8d0
d96a372
15719dc
a274c02
6f35684
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,5 @@ | ||
--- | ||
"stagehand": patch | ||
--- | ||
|
||
Don't close new opened tabs | ||
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -40,7 +40,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 +54,83 @@ 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): | ||
instance._handle_new_page(pw_page) | ||
|
||
context.on("page", handle_page_event) | ||
|
||
return instance | ||
|
||
def _handle_new_page(self, pw_page: Page): | ||
""" | ||
Handle new pages created by the browser (popups, window.open, etc.). | ||
This runs synchronously in the event handler context. | ||
""" | ||
|
||
async def _async_handle(): | ||
try: | ||
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" | ||
) | ||
except Exception as e: | ||
self.stagehand.logger.error( | ||
f"Failed to initialize new page: {str(e)}", category="context" | ||
) | ||
import traceback | ||
|
||
self.stagehand.logger.error( | ||
f"Traceback: {traceback.format_exc()}", category="context" | ||
) | ||
|
||
# Schedule the async work | ||
import asyncio | ||
|
||
|
||
try: | ||
loop = asyncio.get_running_loop() | ||
loop.create_task(_async_handle()) | ||
except RuntimeError: | ||
# No event loop running, which shouldn't happen in normal operation | ||
self.stagehand.logger.error( | ||
"No event loop available to handle new page", category="context" | ||
) | ||
|
||
def __getattr__(self, name): | ||
# Forward attribute lookups to the underlying BrowserContext | ||
attr = getattr(self._context, name) | ||
|
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -33,6 +33,84 @@ | |
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): | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. prevent race condition:
|
||
"""Wait for any pending page switches to complete""" | ||
if hasattr(self._stagehand, "_page_switch_lock"): | ||
async with self._stagehand._page_switch_lock: | ||
pass # Just wait for any ongoing switches | ||
|
||
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 | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. block everything until page is set |
||
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 +244,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 +259,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 +487,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 +507,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 +695,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 | ||
|
Uh oh!
There was an error while loading. Please reload this page.