add timeout when close browser (#1134)

skyvern/webeye/browser_factory.py

@@ -14,7 +14,7 @@ from playwright.async_api import BrowserContext, ConsoleMessage, Error, Page, Pl
 from pydantic import BaseModel, PrivateAttr
 
 from skyvern.config import settings
-from skyvern.constants import REPO_ROOT_DIR
+from skyvern.constants import BROWSER_CLOSE_TIMEOUT, REPO_ROOT_DIR
 from skyvern.exceptions import (
     FailedToNavigateToUrl,
     FailedToReloadPage,
@@ -443,6 +443,8 @@ class BrowserState:
 
     async def close(self, close_browser_on_completion: bool = True) -> None:
         LOG.info("Closing browser state")
+        try:
+            async with asyncio.timeout(BROWSER_CLOSE_TIMEOUT):
                 if self.browser_context and close_browser_on_completion:
                     LOG.info("Closing browser context and its pages")
                     await self.browser_context.close()
@@ -450,10 +452,17 @@ class BrowserState:
                     if self.browser_cleanup is not None:
                         self.browser_cleanup()
                         LOG.info("Main browser cleanup is excuted")
+        except asyncio.TimeoutError:
+            LOG.error("Timeout to close browser context, going to stop playwright directly")
+
+        try:
+            async with asyncio.timeout(BROWSER_CLOSE_TIMEOUT):
                 if self.pw and close_browser_on_completion:
                     LOG.info("Stopping playwright")
                     await self.pw.stop()
                     LOG.info("Playwright is stopped")
+        except asyncio.TimeoutError:
+            LOG.error("Timeout to close playwright, might leave the broswer opening forever")
 
     async def take_screenshot(self, full_page: bool = False, file_path: str | None = None) -> bytes:
         page = await self.__assert_page()

skyvern/webeye/browser_manager.py

@@ -1,12 +1,10 @@
 from __future__ import annotations
 
-import asyncio
 import os
 
 import structlog
 from playwright.async_api import async_playwright
 
-from skyvern.constants import BROWSER_CLOSE_TIMEOUT
 from skyvern.exceptions import MissingBrowserState
 from skyvern.forge.sdk.schemas.tasks import ProxyLocation, Task
 from skyvern.forge.sdk.workflow.models.workflow import WorkflowRun
@@ -208,9 +206,7 @@ class BrowserManager:
     async def cleanup_for_task(self, task_id: str, close_browser_on_completion: bool = True) -> BrowserState | None:
         LOG.info("Cleaning up for task")
         browser_state_to_close = self.pages.pop(task_id, None)
-        try:
         if browser_state_to_close:
-                async with asyncio.timeout(BROWSER_CLOSE_TIMEOUT):
             # Stop tracing before closing the browser if tracing is enabled
             if browser_state_to_close.browser_context and browser_state_to_close.browser_artifacts.traces_dir:
                 trace_path = f"{browser_state_to_close.browser_artifacts.traces_dir}/{task_id}.zip"
@@ -218,8 +214,6 @@ class BrowserManager:
                 LOG.info("Stopped tracing", trace_path=trace_path)
             await browser_state_to_close.close(close_browser_on_completion=close_browser_on_completion)
         LOG.info("Task is cleaned up")
-        except TimeoutError:
-            LOG.warning("Timeout on task cleanup")
 
         return browser_state_to_close