error boundary for streaming

src/browser/custom_context.py

@@ -24,7 +24,7 @@ class CustomBrowserContext(BrowserContext):
         self,
         browser: "Browser",
         config: BrowserContextConfig = BrowserContextConfig(),
-        context: PlaywrightBrowserContext = None,
+        context: PlaywrightBrowserContext | None = None,
     ):
         super(CustomBrowserContext, self).__init__(browser=browser, config=config)
         self.context = context
@@ -34,6 +34,8 @@ class CustomBrowserContext(BrowserContext):
     @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:
@@ -106,6 +108,8 @@ class CustomBrowserContext(BrowserContext):
         """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.")

webui.py

@@ -32,7 +32,6 @@ 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
 
@@ -87,6 +86,7 @@ async def run_browser_agent(
                 window_w=window_w,
                 window_h=window_h,
                 save_recording_path=save_recording_path,
+                save_trace_path=save_trace_path,
                 task=task,
                 max_steps=max_steps,
                 use_vision=use_vision,
@@ -119,7 +119,7 @@ async def run_browser_agent(
 
     finally:
         if browser:
-            await browser.close()
+            browser.close()
 
 async def run_org_agent(
         llm,