Implement Browser context 1

Author: casistack

src/browser/custom_browser.py

@@ -6,22 +6,45 @@
 
 from browser_use.browser.browser import Browser
 from browser_use.browser.context import BrowserContext, BrowserContextConfig
+from playwright.async_api import BrowserContext as PlaywrightBrowserContext
+import logging
 
 from .config import BrowserPersistenceConfig
 from .custom_context import CustomBrowserContext
 
+logger = logging.getLogger(__name__)
 
 class CustomBrowser(Browser):
+    _global_context = None
+
     async def new_context(
         self,
         config: BrowserContextConfig = BrowserContextConfig(),
-        context: CustomBrowserContext = None,
-    ) -> BrowserContext:
-        """Create a browser context"""
+        context: PlaywrightBrowserContext = None,
+    ) -> 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"""
-        # Check if persistence is enabled before closing
         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()
+        else:
+            logger.info("Skipping browser close due to persistent session")

src/browser/custom_context.py

@@ -12,6 +12,7 @@
 from browser_use.browser.browser import Browser
 from browser_use.browser.context import BrowserContext, BrowserContextConfig
 from playwright.async_api import Browser as PlaywrightBrowser
+from playwright.async_api import BrowserContext as PlaywrightBrowserContext
 
 from .config import BrowserPersistenceConfig
 logger = logging.getLogger(__name__)
@@ -22,19 +23,21 @@ def __init__(
         self,
         browser: "Browser",
         config: BrowserContextConfig = BrowserContextConfig(),
-        context: BrowserContext = None,
+        context: PlaywrightBrowserContext = None,
     ):
         super(CustomBrowserContext, self).__init__(browser=browser, config=config)
         self.context = context
-        
+        self._persistence_config = BrowserPersistenceConfig.from_env()
 
-    async def _create_context(self, browser: PlaywrightBrowser):
+    async def _create_context(self, browser: PlaywrightBrowser) -> PlaywrightBrowserContext:
         """Creates a new browser context with anti-detection measures and loads cookies if available."""
         # If we have a context, return it directly
         if self.context:
             return self.context
-        if self.browser.config.chrome_instance_path and len(browser.contexts) > 0:
-            # Connect to existing Chrome instance instead of creating new one
+
+        # 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")
             context = browser.contexts[0]
         else:
             # Original code for creating new context
@@ -49,7 +52,7 @@ async def _create_context(self, browser: PlaywrightBrowser):
                 bypass_csp=self.config.disable_security,
                 ignore_https_errors=self.config.disable_security,
                 record_video_dir=self.config.save_recording_path,
-                record_video_size=self.config.browser_window_size,  # set record video size, same as windows size
+                record_video_size=self.config.browser_window_size,
             )
 
         if self.config.trace_path:
@@ -96,3 +99,8 @@ async def _create_context(self, browser: PlaywrightBrowser):
         )
 
         return context
+
+    async def close(self):
+        """Override close to respect persistence setting"""
+        if not self._persistence_config.persistent_session:
+            await super().close()

webui.py

@@ -35,6 +35,16 @@
 from src.controller.custom_controller import CustomController
 from src.utils import utils
 from src.utils.utils import update_model_dropdown
+from src.browser.config import BrowserPersistenceConfig
+from src.browser.custom_browser import CustomBrowser
+from src.browser.custom_context import CustomBrowserContext
+from browser_use.browser.browser import BrowserConfig
+from browser_use.browser.context import BrowserContextConfig, BrowserContextWindowSize
+
+# Global variables for persistence
+_global_browser = None
+_global_browser_context = None
+_global_playwright = None
 
 async def run_browser_agent(
         agent_type,
@@ -196,96 +206,107 @@ async def run_custom_agent(
         max_actions_per_step,
         tool_call_in_content
 ):
+    global _global_browser, _global_browser_context, _global_playwright
+    
     controller = CustomController()
-    playwright = None
-    browser_context_ = None
+    persistence_config = BrowserPersistenceConfig.from_env()
+    
     try:
-        if use_own_browser:
-            playwright = await async_playwright().start()
-            chrome_exe = os.getenv("CHROME_PATH", "")
-            chrome_use_data = os.getenv("CHROME_USER_DATA", "")
-
-            if chrome_exe == "":
-                chrome_exe = None
-            elif not os.path.exists(chrome_exe):
-                raise ValueError(f"Chrome executable not found at {chrome_exe}")
-
-            if chrome_use_data == "":
-                chrome_use_data = None
-
-            browser_context_ = await playwright.chromium.launch_persistent_context(
-                user_data_dir=chrome_use_data,
-                executable_path=chrome_exe,
-                no_viewport=False,
-                headless=headless,  # 保持浏览器窗口可见
-                user_agent=(
-                    "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 "
-                    "(KHTML, like Gecko) Chrome/85.0.4183.102 Safari/537.36"
-                ),
-                java_script_enabled=True,
-                bypass_csp=disable_security,
-                ignore_https_errors=disable_security,
-                record_video_dir=save_recording_path if save_recording_path else None,
-                record_video_size={"width": window_w, "height": window_h},
+        # Initialize global browser if needed
+        if _global_browser is None:
+            _global_browser = CustomBrowser(
+                config=BrowserConfig(
+                    headless=headless,
+                    disable_security=disable_security,
+                    extra_chromium_args=[f"--window-size={window_w},{window_h}"],
+                )
             )
-        else:
-            browser_context_ = None
 
-        browser = CustomBrowser(
-            config=BrowserConfig(
-                headless=headless,
-                disable_security=disable_security,
-                extra_chromium_args=[f"--window-size={window_w},{window_h}"],
-            )
-        )
-        async with await browser.new_context(
-                config=BrowserContextConfig(
-                    trace_path=save_trace_path if save_trace_path else None,
-                    save_recording_path=save_recording_path
-                    if save_recording_path
-                    else None,
+        # Handle browser context based on configuration
+        if use_own_browser:
+            if _global_browser_context is None:
+                _global_playwright = await async_playwright().start()
+                chrome_exe = os.getenv("CHROME_PATH", "")
+                chrome_use_data = os.getenv("CHROME_USER_DATA", "")
+
+                browser_context = await _global_playwright.chromium.launch_persistent_context(
+                    user_data_dir=chrome_use_data,
+                    executable_path=chrome_exe,
                     no_viewport=False,
-                    browser_window_size=BrowserContextWindowSize(
-                        width=window_w, height=window_h
+                    headless=headless,
+                    user_agent=(
+                        "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 "
+                        "(KHTML, like Gecko) Chrome/85.0.4183.102 Safari/537.36"
                     ),
-                ),
-                context=browser_context_,
-        ) as browser_context:
-            agent = CustomAgent(
-                task=task,
-                add_infos=add_infos,
-                use_vision=use_vision,
-                llm=llm,
-                browser_context=browser_context,
-                controller=controller,
-                system_prompt_class=CustomSystemPrompt,
-                max_actions_per_step=max_actions_per_step,
-                tool_call_in_content=tool_call_in_content
-            )
-            history = await agent.run(max_steps=max_steps)
+                    java_script_enabled=True,
+                    bypass_csp=disable_security,
+                    ignore_https_errors=disable_security,
+                    record_video_dir=save_recording_path if save_recording_path else None,
+                    record_video_size={"width": window_w, "height": window_h},
+                )
+                _global_browser_context = await _global_browser.new_context(
+                    config=BrowserContextConfig(
+                        trace_path=save_trace_path if save_trace_path else None,
+                        save_recording_path=save_recording_path if save_recording_path else None,
+                        no_viewport=False,
+                        browser_window_size=BrowserContextWindowSize(
+                            width=window_w, height=window_h
+                        ),
+                    ),
+                    context=browser_context,
+                )
+        else:
+            if _global_browser_context is None:
+                _global_browser_context = await _global_browser.new_context(
+                    config=BrowserContextConfig(
+                        trace_path=save_trace_path if save_trace_path else None,
+                        save_recording_path=save_recording_path if save_recording_path else None,
+                        no_viewport=False,
+                        browser_window_size=BrowserContextWindowSize(
+                            width=window_w, height=window_h
+                        ),
+                    ),
+                )
 
-            final_result = history.final_result()
-            errors = history.errors()
-            model_actions = history.model_actions()
-            model_thoughts = history.model_thoughts()
+        # Create and run agent
+        agent = CustomAgent(
+            task=task,
+            add_infos=add_infos,
+            use_vision=use_vision,
+            llm=llm,
+            browser_context=_global_browser_context,
+            controller=controller,
+            system_prompt_class=CustomSystemPrompt,
+            max_actions_per_step=max_actions_per_step,
+            tool_call_in_content=tool_call_in_content
+        )
+        history = await agent.run(max_steps=max_steps)
+
+        final_result = history.final_result()
+        errors = history.errors()
+        model_actions = history.model_actions()
+        model_thoughts = history.model_thoughts()
 
     except Exception as e:
         import traceback
-
         traceback.print_exc()
-        final_result = ""
         errors = str(e) + "\n" + traceback.format_exc()
-        model_actions = ""
-        model_thoughts = ""
+        
     finally:
-        # 显式关闭持久化上下文
-        if browser_context_:
-            await browser_context_.close()
-
-        # 关闭 Playwright 对象
-        if playwright:
-            await playwright.stop()
-        await browser.close()
+        # Handle cleanup based on persistence configuration
+        if not persistence_config.persistent_session:
+            if _global_browser_context:
+                await _global_browser_context.close()
+                _global_browser_context = None
+
+            if _global_playwright:
+                await _global_playwright.stop()
+                _global_playwright = None
+
+            if _global_browser:
+                await _global_browser.close()
+                _global_browser = None
+
     return final_result, errors, model_actions, model_thoughts
 
 # Define the theme map globally