diff --git a/pyproject.toml b/pyproject.toml index b3c6ab92..a89a686d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -38,6 +38,7 @@ dependencies = [ "markdownify>=1.0.0,<2.0.0", "requests>=2.28.0,<3.0.0", "aiohttp>=3.8.0,<4.0.0", + "typing_extensions>=4.0.0,<5.0.0", # Note: Always want the latest tzdata "tzdata ; platform_system == 'Windows'", "botocore>=1.39.7,<2.0.0", diff --git a/src/strands_tools/browser/agent_core_browser.py b/src/strands_tools/browser/agent_core_browser.py index c6dea386..24a7fd14 100644 --- a/src/strands_tools/browser/agent_core_browser.py +++ b/src/strands_tools/browser/agent_core_browser.py @@ -10,6 +10,7 @@ from bedrock_agentcore.tools.browser_client import BrowserClient as AgentCoreBrowserClient from playwright.async_api import Browser as PlaywrightBrowser +from typing_extensions import override from ..utils.aws_util import resolve_region from .browser import Browser @@ -59,6 +60,34 @@ async def create_browser_session(self) -> PlaywrightBrowser: return browser + @override + async def _setup_session_from_browser(self, browser_or_context): + """Setup session for AgentCoreBrowser using existing CDP context. + + AgentCoreBrowser connects via CDP and uses the default context + that comes with the connection, avoiding the need to create a new one. + + Args: + browser_or_context: PlaywrightBrowser from CDP connection + + Returns: + Tuple of (browser, context, page) + """ + session_browser = browser_or_context + + # CDP connections should have a default context + if not session_browser.contexts: + raise RuntimeError( + "AgentCoreBrowser CDP connection has no contexts. " + "This may indicate a connection issue with the remote browser." + ) + + # Use the existing default context from CDP connection + session_context = session_browser.contexts[0] + session_page = await session_context.new_page() + + return session_browser, session_context, session_page + def close_platform(self) -> None: for client in self._client_dict.values(): try: diff --git a/src/strands_tools/browser/browser.py b/src/strands_tools/browser/browser.py index 7a706b24..e2e314c3 100644 --- a/src/strands_tools/browser/browser.py +++ b/src/strands_tools/browser/browser.py @@ -221,6 +221,31 @@ async def create_browser_session(self) -> PlaywrightBrowser: """ ... + async def _setup_session_from_browser(self, browser_or_context): + """Setup session components from browser or context. + + This method can be overridden by subclasses to customize how + browser, context, and page are extracted from the session object. + + Args: + browser_or_context: The object returned by create_browser_session() + + Returns: + Tuple of (browser, context, page) + """ + if isinstance(browser_or_context, PlaywrightBrowser): + # Normal non-persistent case + session_browser = browser_or_context + session_context = await session_browser.new_context() + session_page = await session_context.new_page() + else: + # Persistent context case + session_context = browser_or_context + session_browser = session_context.browser + session_page = await session_context.new_page() + + return session_browser, session_context, session_page + # Session Management Methods def init_session(self, action: InitSessionAction) -> Dict[str, Any]: """Initialize a new browser session.""" @@ -238,19 +263,10 @@ async def _async_init_session(self, action: InitSessionAction) -> Dict[str, Any] try: # Create new browser instance for this session - session = await self.create_browser_session() + browser_or_context = await self.create_browser_session() - if isinstance(session, PlaywrightBrowser): - # Normal non-persistent case - session_browser = session - session_context = await session_browser.new_context() - session_page = await session_context.new_page() - - else: - # Persistent context case - session_context = session - session_browser = session_context.browser - session_page = await session_context.new_page() + # Let subclasses customize how to extract browser, context, and page + session_browser, session_context, session_page = await self._setup_session_from_browser(browser_or_context) # Create and store session object session = BrowserSession( diff --git a/tests/test_editor.py b/tests/test_editor.py index 9d340d4b..cdc61595 100644 --- a/tests/test_editor.py +++ b/tests/test_editor.py @@ -259,12 +259,7 @@ def test_str_replace_no_backup(self, mock_user_input, temp_file, clean_content_h """Test str_replace without creating backup when EDITOR_DISABLE_BACKUP is set.""" mock_user_input.return_value = "y" - result = editor.editor( - command="str_replace", - path=temp_file, - old_str="Line 2", - new_str="Modified Line 2" - ) + result = editor.editor(command="str_replace", path=temp_file, old_str="Line 2", new_str="Modified Line 2") assert result["status"] == "success" @@ -278,12 +273,7 @@ def test_pattern_replace_no_backup(self, mock_user_input, temp_file, clean_conte """Test pattern_replace without creating backup.""" mock_user_input.return_value = "y" - result = editor.editor( - command="pattern_replace", - path=temp_file, - pattern="Line.*", - new_str="Updated Line" - ) + result = editor.editor(command="pattern_replace", path=temp_file, pattern="Line.*", new_str="Updated Line") assert result["status"] == "success" backup_path = f"{temp_file}.bak" @@ -295,12 +285,7 @@ def test_insert_no_backup(self, mock_user_input, temp_file, clean_content_histor """Test insert without creating backup.""" mock_user_input.return_value = "y" - result = editor.editor( - command="insert", - path=temp_file, - new_str="New line", - insert_line=2 - ) + result = editor.editor(command="insert", path=temp_file, new_str="New line", insert_line=2) assert result["status"] == "success" backup_path = f"{temp_file}.bak" @@ -315,12 +300,7 @@ def test_backup_created_by_default(self, mock_user_input, temp_file, clean_conte mock_user_input.return_value = "y" - result = editor.editor( - command="str_replace", - path=temp_file, - old_str="Line 2", - new_str="Modified Line 2" - ) + result = editor.editor(command="str_replace", path=temp_file, old_str="Line 2", new_str="Modified Line 2") assert result["status"] == "success" backup_path = f"{temp_file}.bak"