mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-22 14:28:55 +00:00
e16606672a
git-subtree-dir: packages/tools git-subtree-split: 78317b9c127f18bd040c1d77e3c0840cdc9a5b38
260 lines
8.4 KiB
Python
260 lines
8.4 KiB
Python
from __future__ import annotations
|
|
|
|
import logging
|
|
from typing import TYPE_CHECKING, Dict, Tuple
|
|
|
|
if TYPE_CHECKING:
|
|
from playwright.async_api import Browser as AsyncBrowser
|
|
from playwright.sync_api import Browser as SyncBrowser
|
|
from bedrock_agentcore.tools.browser_client import BrowserClient
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class BrowserSessionManager:
|
|
"""
|
|
Manages browser sessions for different threads.
|
|
|
|
This class maintains separate browser sessions for different threads,
|
|
enabling concurrent usage of browsers in multi-threaded environments.
|
|
Browsers are created lazily only when needed by tools.
|
|
"""
|
|
|
|
def __init__(self, region: str = "us-west-2"):
|
|
"""
|
|
Initialize the browser session manager.
|
|
|
|
Args:
|
|
region: AWS region for browser client
|
|
"""
|
|
self.region = region
|
|
self._async_sessions: Dict[str, Tuple[BrowserClient, AsyncBrowser]] = {}
|
|
self._sync_sessions: Dict[str, Tuple[BrowserClient, SyncBrowser]] = {}
|
|
|
|
async def get_async_browser(self, thread_id: str) -> AsyncBrowser:
|
|
"""
|
|
Get or create an async browser for the specified thread.
|
|
|
|
Args:
|
|
thread_id: Unique identifier for the thread requesting the browser
|
|
|
|
Returns:
|
|
An async browser instance specific to the thread
|
|
"""
|
|
if thread_id in self._async_sessions:
|
|
return self._async_sessions[thread_id][1]
|
|
|
|
return await self._create_async_browser_session(thread_id)
|
|
|
|
def get_sync_browser(self, thread_id: str) -> SyncBrowser:
|
|
"""
|
|
Get or create a sync browser for the specified thread.
|
|
|
|
Args:
|
|
thread_id: Unique identifier for the thread requesting the browser
|
|
|
|
Returns:
|
|
A sync browser instance specific to the thread
|
|
"""
|
|
if thread_id in self._sync_sessions:
|
|
return self._sync_sessions[thread_id][1]
|
|
|
|
return self._create_sync_browser_session(thread_id)
|
|
|
|
async def _create_async_browser_session(self, thread_id: str) -> AsyncBrowser:
|
|
"""
|
|
Create a new async browser session for the specified thread.
|
|
|
|
Args:
|
|
thread_id: Unique identifier for the thread
|
|
|
|
Returns:
|
|
The newly created async browser instance
|
|
|
|
Raises:
|
|
Exception: If browser session creation fails
|
|
"""
|
|
from bedrock_agentcore.tools.browser_client import BrowserClient
|
|
browser_client = BrowserClient(region=self.region)
|
|
|
|
try:
|
|
# Start browser session
|
|
browser_client.start()
|
|
|
|
# Get WebSocket connection info
|
|
ws_url, headers = browser_client.generate_ws_headers()
|
|
|
|
logger.info(
|
|
f"Connecting to async WebSocket endpoint for thread {thread_id}: {ws_url}"
|
|
)
|
|
|
|
from playwright.async_api import async_playwright
|
|
|
|
# Connect to browser using Playwright
|
|
playwright = await async_playwright().start()
|
|
browser = await playwright.chromium.connect_over_cdp(
|
|
endpoint_url=ws_url, headers=headers, timeout=30000
|
|
)
|
|
logger.info(
|
|
f"Successfully connected to async browser for thread {thread_id}"
|
|
)
|
|
|
|
# Store session resources
|
|
self._async_sessions[thread_id] = (browser_client, browser)
|
|
|
|
return browser
|
|
|
|
except Exception as e:
|
|
logger.error(
|
|
f"Failed to create async browser session for thread {thread_id}: {e}"
|
|
)
|
|
|
|
# Clean up resources if session creation fails
|
|
if browser_client:
|
|
try:
|
|
browser_client.stop()
|
|
except Exception as cleanup_error:
|
|
logger.warning(f"Error cleaning up browser client: {cleanup_error}")
|
|
|
|
raise
|
|
|
|
def _create_sync_browser_session(self, thread_id: str) -> SyncBrowser:
|
|
"""
|
|
Create a new sync browser session for the specified thread.
|
|
|
|
Args:
|
|
thread_id: Unique identifier for the thread
|
|
|
|
Returns:
|
|
The newly created sync browser instance
|
|
|
|
Raises:
|
|
Exception: If browser session creation fails
|
|
"""
|
|
from bedrock_agentcore.tools.browser_client import BrowserClient
|
|
browser_client = BrowserClient(region=self.region)
|
|
|
|
try:
|
|
# Start browser session
|
|
browser_client.start()
|
|
|
|
# Get WebSocket connection info
|
|
ws_url, headers = browser_client.generate_ws_headers()
|
|
|
|
logger.info(
|
|
f"Connecting to sync WebSocket endpoint for thread {thread_id}: {ws_url}"
|
|
)
|
|
|
|
from playwright.sync_api import sync_playwright
|
|
|
|
# Connect to browser using Playwright
|
|
playwright = sync_playwright().start()
|
|
browser = playwright.chromium.connect_over_cdp(
|
|
endpoint_url=ws_url, headers=headers, timeout=30000
|
|
)
|
|
logger.info(
|
|
f"Successfully connected to sync browser for thread {thread_id}"
|
|
)
|
|
|
|
# Store session resources
|
|
self._sync_sessions[thread_id] = (browser_client, browser)
|
|
|
|
return browser
|
|
|
|
except Exception as e:
|
|
logger.error(
|
|
f"Failed to create sync browser session for thread {thread_id}: {e}"
|
|
)
|
|
|
|
# Clean up resources if session creation fails
|
|
if browser_client:
|
|
try:
|
|
browser_client.stop()
|
|
except Exception as cleanup_error:
|
|
logger.warning(f"Error cleaning up browser client: {cleanup_error}")
|
|
|
|
raise
|
|
|
|
async def close_async_browser(self, thread_id: str) -> None:
|
|
"""
|
|
Close the async browser session for the specified thread.
|
|
|
|
Args:
|
|
thread_id: Unique identifier for the thread
|
|
"""
|
|
if thread_id not in self._async_sessions:
|
|
logger.warning(f"No async browser session found for thread {thread_id}")
|
|
return
|
|
|
|
browser_client, browser = self._async_sessions[thread_id]
|
|
|
|
# Close browser
|
|
if browser:
|
|
try:
|
|
await browser.close()
|
|
except Exception as e:
|
|
logger.warning(
|
|
f"Error closing async browser for thread {thread_id}: {e}"
|
|
)
|
|
|
|
# Stop browser client
|
|
if browser_client:
|
|
try:
|
|
browser_client.stop()
|
|
except Exception as e:
|
|
logger.warning(
|
|
f"Error stopping browser client for thread {thread_id}: {e}"
|
|
)
|
|
|
|
# Remove session from dictionary
|
|
del self._async_sessions[thread_id]
|
|
logger.info(f"Async browser session cleaned up for thread {thread_id}")
|
|
|
|
def close_sync_browser(self, thread_id: str) -> None:
|
|
"""
|
|
Close the sync browser session for the specified thread.
|
|
|
|
Args:
|
|
thread_id: Unique identifier for the thread
|
|
"""
|
|
if thread_id not in self._sync_sessions:
|
|
logger.warning(f"No sync browser session found for thread {thread_id}")
|
|
return
|
|
|
|
browser_client, browser = self._sync_sessions[thread_id]
|
|
|
|
# Close browser
|
|
if browser:
|
|
try:
|
|
browser.close()
|
|
except Exception as e:
|
|
logger.warning(
|
|
f"Error closing sync browser for thread {thread_id}: {e}"
|
|
)
|
|
|
|
# Stop browser client
|
|
if browser_client:
|
|
try:
|
|
browser_client.stop()
|
|
except Exception as e:
|
|
logger.warning(
|
|
f"Error stopping browser client for thread {thread_id}: {e}"
|
|
)
|
|
|
|
# Remove session from dictionary
|
|
del self._sync_sessions[thread_id]
|
|
logger.info(f"Sync browser session cleaned up for thread {thread_id}")
|
|
|
|
async def close_all_browsers(self) -> None:
|
|
"""Close all browser sessions."""
|
|
# Close all async browsers
|
|
async_thread_ids = list(self._async_sessions.keys())
|
|
for thread_id in async_thread_ids:
|
|
await self.close_async_browser(thread_id)
|
|
|
|
# Close all sync browsers
|
|
sync_thread_ids = list(self._sync_sessions.keys())
|
|
for thread_id in sync_thread_ids:
|
|
self.close_sync_browser(thread_id)
|
|
|
|
logger.info("All browser sessions closed") |