new stream function without the need to modify custom context

Author: katiue

src/browser/custom_browser.py

@@ -4,6 +4,16 @@
 # @ProjectName: browser-use-webui
 # @FileName: browser.py
 
+import asyncio
+
+from playwright.async_api import Browser as PlaywrightBrowser
+from playwright.async_api import (
+	BrowserContext as PlaywrightBrowserContext,
+)
+from playwright.async_api import (
+	Playwright,
+	async_playwright,
+)
 from browser_use.browser.browser import Browser
 from browser_use.browser.context import BrowserContext, BrowserContextConfig
 from playwright.async_api import BrowserContext as PlaywrightBrowserContext
@@ -15,36 +25,102 @@
 logger = logging.getLogger(__name__)
 
 class CustomBrowser(Browser):
-    _global_context = None
 
     async def new_context(
         self,
-        config: BrowserContextConfig = BrowserContextConfig(),
-        context: PlaywrightBrowserContext = None,
+        config: BrowserContextConfig = BrowserContextConfig()
     ) -> CustomBrowserContext:
-        """Create a browser context with persistence support"""
-        persistence_config = BrowserPersistenceConfig.from_env()
-        
-        if persistence_config.persistent_session:
-            if CustomBrowser._global_context is not None:
-                logger.info("Reusing existing persistent browser context")
-                return CustomBrowser._global_context
-            
-            context_instance = CustomBrowserContext(config=config, browser=self, context=context)
-            CustomBrowser._global_context = context_instance
-            logger.info("Created new persistent browser context")
-            return context_instance
-        
-        logger.info("Creating non-persistent browser context")
-        return CustomBrowserContext(config=config, browser=self, context=context)
-
-    async def close(self):
-        """Override close to respect persistence setting"""
-        persistence_config = BrowserPersistenceConfig.from_env()
-        if not persistence_config.persistent_session:
-            if CustomBrowser._global_context is not None:
-                await CustomBrowser._global_context.close()
-                CustomBrowser._global_context = None
-            await super().close()
+        return CustomBrowserContext(config=config, browser=self)
+
+    async def _setup_browser(self, playwright: Playwright) -> PlaywrightBrowser:
+        """Sets up and returns a Playwright Browser instance with anti-detection measures."""
+        if self.config.wss_url:
+            browser = await playwright.chromium.connect(self.config.wss_url)
+            return browser
+        elif self.config.chrome_instance_path:
+            import subprocess
+
+            import requests
+
+            try:
+                # Check if browser is already running
+                response = requests.get('http://localhost:9222/json/version', timeout=2)
+                if response.status_code == 200:
+                    logger.info('Reusing existing Chrome instance')
+                    browser = await playwright.chromium.connect_over_cdp(
+                        endpoint_url='http://localhost:9222',
+                        timeout=20000,  # 20 second timeout for connection
+                    )
+                    return browser
+            except requests.ConnectionError:
+                logger.debug('No existing Chrome instance found, starting a new one')
+
+            # Start a new Chrome instance
+            subprocess.Popen(
+                [
+                    self.config.chrome_instance_path,
+                    '--remote-debugging-port=9222',
+                ],
+                stdout=subprocess.DEVNULL,
+                stderr=subprocess.DEVNULL,
+            )
+
+            # Attempt to connect again after starting a new instance
+            for _ in range(10):
+                try:
+                    response = requests.get('http://localhost:9222/json/version', timeout=2)
+                    if response.status_code == 200:
+                        break
+                except requests.ConnectionError:
+                    pass
+                await asyncio.sleep(1)
+
+            try:
+                browser = await playwright.chromium.connect_over_cdp(
+                    endpoint_url='http://localhost:9222',
+                    timeout=20000,  # 20 second timeout for connection
+                )
+                return browser
+            except Exception as e:
+                logger.error(f'Failed to start a new Chrome instance.: {str(e)}')
+                raise RuntimeError(
+                    ' To start chrome in Debug mode, you need to close all existing Chrome instances and try again otherwise we can not connect to the instance.'
+                )
+
         else:
-            logger.info("Skipping browser close due to persistent session")
+            try:
+                disable_security_args = []
+                if self.config.disable_security:
+                    disable_security_args = [
+                        '--disable-web-security',
+                        '--disable-site-isolation-trials',
+                        '--disable-features=IsolateOrigins,site-per-process',
+                    ]
+
+                browser = await playwright.chromium.launch(
+                    headless=self.config.headless,
+                    args=[
+                             '--no-sandbox',
+                             '--disable-blink-features=AutomationControlled',
+                             '--disable-infobars',
+                             '--disable-background-timer-throttling',
+                             '--disable-popup-blocking',
+                             '--disable-backgrounding-occluded-windows',
+                             '--disable-renderer-backgrounding',
+                             '--disable-window-activation',
+                             '--disable-focus-on-load',
+                             '--no-first-run',
+                             '--no-default-browser-check',
+                             '--no-startup-window',
+                             '--window-position=0,0',
+                             # '--window-size=1280,1000',
+                         ]
+                         + disable_security_args
+                         + self.config.extra_chromium_args,
+                    proxy=self.config.proxy,
+                )
+
+                return browser
+            except Exception as e:
+                logger.error(f'Failed to initialize Playwright browser: {str(e)}')
+                raise

src/browser/custom_context.py

@@ -15,41 +15,28 @@
 from playwright.async_api import BrowserContext as PlaywrightBrowserContext
 
 from .config import BrowserPersistenceConfig
-
 logger = logging.getLogger(__name__)
 
 
 class CustomBrowserContext(BrowserContext):
     def __init__(
         self,
         browser: "Browser",
-        config: BrowserContextConfig = BrowserContextConfig(),
-        context: PlaywrightBrowserContext = None,
+        config: BrowserContextConfig = BrowserContextConfig()
     ):
         super(CustomBrowserContext, self).__init__(browser=browser, config=config)
-        self.context = context
-        self._page = None
-        self._persistence_config = BrowserPersistenceConfig.from_env()
-
-    @property
-    def impl_context(self) -> PlaywrightBrowserContext:
-        """Returns the underlying Playwright context implementation"""
-        if self.context is None:
-            raise RuntimeError("Failed to create or retrieve a browser context.")
-        return self.context
 
     async def _create_context(self, browser: PlaywrightBrowser) -> PlaywrightBrowserContext:
         """Creates a new browser context with anti-detection measures and loads cookies if available."""
-        if self.context:
-            return self.context
+        # If we have a context, return it directly
 
         # Check if we should use existing context for persistence
-        if self._persistence_config.persistent_session and len(browser.contexts) > 0:
-            logger.info("Using existing persistent context.")
-            self.context = browser.contexts[0]
+        if self.browser.config.chrome_instance_path and len(browser.contexts) > 0:
+            # Connect to existing Chrome instance instead of creating new one
+            context = browser.contexts[0]
         else:
-            logger.info("Creating a new browser context.")
-            self.context = await browser.new_context(
+            # Original code for creating new context
+            context = await browser.new_context(
                 viewport=self.config.browser_window_size,
                 no_viewport=False,
                 user_agent=(
@@ -63,19 +50,20 @@ async def _create_context(self, browser: PlaywrightBrowser) -> PlaywrightBrowser
                 record_video_size=self.config.browser_window_size,
             )
 
-        # Handle tracing
         if self.config.trace_path:
-            await self.context.tracing.start(screenshots=True, snapshots=True, sources=True)
+            await context.tracing.start(screenshots=True, snapshots=True, sources=True)
 
         # Load cookies if they exist
         if self.config.cookies_file and os.path.exists(self.config.cookies_file):
             with open(self.config.cookies_file, "r") as f:
                 cookies = json.load(f)
-                logger.info(f"Loaded {len(cookies)} cookies from {self.config.cookies_file}.")
-                await self.context.add_cookies(cookies)
+                logger.info(
+                    f"Loaded {len(cookies)} cookies from {self.config.cookies_file}"
+                )
+                await context.add_cookies(cookies)
 
         # Expose anti-detection scripts
-        await self.context.add_init_script(
+        await context.add_init_script(
             """
             // Webdriver property
             Object.defineProperty(navigator, 'webdriver', {
@@ -105,41 +93,4 @@ async def _create_context(self, browser: PlaywrightBrowser) -> PlaywrightBrowser
             """
         )
 
-        # Create initial page if none exists
-        if not self.context.pages:
-            self._page = await self.context.new_page()
-            await self._page.goto('about:blank')
-
-        return self.context
-
-    async def new_page(self):
-        """Creates and returns a new page in this context."""
-        if not self.context:
-            await self._create_context(await self.browser.get_playwright_browser())
-        return await self.context.new_page()
-
-    async def get_current_page(self):
-        """Returns the current page or creates one if none exists."""
-        if not self.context:
-            await self._create_context(await self.browser.get_playwright_browser())
-        if not self.context:
-            raise RuntimeError("Browser context is not initialized.")
-        pages = self.context.pages
-        if not pages:
-            logger.warning("No existing pages in the context. Creating a new page.")
-            return await self.context.new_page()
-        return pages[0]
-
-    async def close(self):
-        """Override close to respect persistence setting."""
-        if not self._persistence_config.persistent_session and self.context:
-            await self.context.close()
-            self.context = None
-
-    @property
-    def pages(self):
-        """Returns list of pages in the context."""
-        if not self.context:
-            logger.warning("Attempting to access pages but context is not initialized.")
-            return []
-        return self.context.pages
+        return context

src/utils/utils.py

@@ -15,6 +15,7 @@
 from langchain_ollama import ChatOllama
 from langchain_openai import AzureChatOpenAI, ChatOpenAI
 import gradio as gr
+from src.browser.custom_context import CustomBrowserContext
 
 def get_llm_model(provider: str, **kwargs):
     """
@@ -164,36 +165,30 @@ def get_latest_files(directory: str, file_types: list = ['.webm', '.zip']) -> Di
             print(f"Error getting latest {file_type} file: {e}")
 
     return latest_files
-
-async def capture_screenshot(browser_context) -> str:
+async def capture_screenshot(browser_context: CustomBrowserContext) -> str:
     """Capture and encode a screenshot"""
+    latest_screenshot = ""
     try:
-        # Get the implementation context - handle both direct Playwright context and wrapped context
-        context = browser_context
-        if hasattr(browser_context, 'context'):
-            context = browser_context.context
-        
-        if not context:
-            return "<div>No browser context available</div>"
-            
-        # Get all pages
-        pages = context.pages
-        if not pages:
-            return "<div>Waiting for page to be available...</div>"
-
-        # Use the first non-blank page or fallback to first page
-        active_page = None
-        for page in pages:
-            if page.url != 'about:blank':
-                active_page = page
-                break
-        
-        if not active_page and pages:
-            active_page = pages[0]
-            
-        if not active_page:
-            return "<div>No active page available</div>"
+        # Extract the Playwright browser instance
+        playwright_browser = browser_context.browser.playwright_browser  # Ensure this is correct.
 
+        # Check if the browser instance is valid and if an existing context can be reused
+        if playwright_browser and playwright_browser.contexts:
+            playwright_context = playwright_browser.contexts[0]
+        else:
+            return latest_screenshot
+
+        # Access pages in the context
+        if playwright_context:
+            pages = playwright_context.pages
+
+        # Use an existing page or create a new one if none exist
+        if pages:
+            active_page = pages[0]
+            for page in pages:
+                if page.url != "about:blank":
+                    active_page = page
+        
         # Take screenshot
         try:
             screenshot = await active_page.screenshot(
@@ -202,9 +197,9 @@ async def capture_screenshot(browser_context) -> str:
                 scale="css"
             )
             encoded = base64.b64encode(screenshot).decode('utf-8')
-            return f'<img src="data:image/jpeg;base64,{encoded}" style="width:100%; max-width:1200px; border:1px solid #ccc;">'
+            return f'<img src="data:image/jpeg;base64,{encoded}" style="width:80vw; height:90vh ; border:1px solid #ccc;">'
         except Exception as e:
-            return f"<div class='error'>Screenshot failed: {str(e)}</div>"
-            
+            return f"<div class='error' style='width:80vw; height:90vh'>Screenshot failed: {str(e)}</div>"
+
     except Exception as e:
-        return f"<div class='error'>Screenshot error: {str(e)}</div>"
+        return f"<div class='error' style='width:80vw; height:90vh'>Screenshot error: {str(e)}</div>"