Add persist_browser_session flag to workflows (#777)

skyvern/webeye/browser_factory.py

@@ -22,6 +22,7 @@ from skyvern.exceptions import (
     UnknownBrowserType,
     UnknownErrorWhileCreatingBrowserContext,
 )
+from skyvern.forge import app
 from skyvern.forge.sdk.core.skyvern_context import current
 from skyvern.forge.sdk.schemas.tasks import ProxyLocation
 from skyvern.forge.sdk.settings_manager import SettingsManager
@@ -89,11 +90,13 @@ class BrowserContextFactory:
         video_artifacts: list[VideoArtifact] | None = None,
         har_path: str | None = None,
         traces_dir: str | None = None,
+        browser_session_dir: str | None = None,
     ) -> BrowserArtifacts:
         return BrowserArtifacts(
             video_artifacts=video_artifacts or [],
             har_path=har_path,
             traces_dir=traces_dir,
+            browser_session_dir=browser_session_dir,
         )
 
     @classmethod
@@ -137,6 +140,7 @@ class BrowserArtifacts(BaseModel):
     video_artifacts: list[VideoArtifact] = []
     har_path: str | None = None
     traces_dir: str | None = None
+    browser_session_dir: str | None = None
 
 
 async def _create_headless_chromium(
@@ -386,3 +390,9 @@ class BrowserState:
     async def take_screenshot(self, full_page: bool = False, file_path: str | None = None) -> bytes:
         page = await self.__assert_page()
         return await SkyvernFrame.take_screenshot(page=page, full_page=full_page, file_path=file_path)
+
+    async def store_browser_session(self, organization_id: str, workflow_permanent_id: str) -> None:
+        if self.browser_artifacts.browser_session_dir:
+            await app.STORAGE.store_browser_session(
+                organization_id, workflow_permanent_id, self.browser_artifacts.browser_session_dir
+            )

skyvern/webeye/browser_manager.py

@@ -9,7 +9,7 @@ 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
+from skyvern.forge.sdk.workflow.models.workflow import Workflow, WorkflowRun
 from skyvern.webeye.browser_factory import BrowserContextFactory, BrowserState, VideoArtifact
 
 LOG = structlog.get_logger()
@@ -182,6 +182,7 @@ class BrowserManager:
 
     async def cleanup_for_workflow_run(
         self,
+        workflow: Workflow,
         workflow_run_id: str,
         task_ids: list[str],
         close_browser_on_completion: bool = True,
@@ -195,6 +196,13 @@ class BrowserManager:
                 await browser_state_to_close.browser_context.tracing.stop(path=trace_path)
                 LOG.info("Stopped tracing", trace_path=trace_path)
 
+            if workflow.persist_browser_session:
+                await browser_state_to_close.store_browser_session(
+                    organization_id=workflow.organization_id,
+                    workflow_permanent_id=workflow.workflow_permanent_id,
+                )
+                LOG.info("Persisted browser session for workflow run", workflow_run_id=workflow_run_id)
+
             await browser_state_to_close.close(close_browser_on_completion=close_browser_on_completion)
         for task_id in task_ids:
             self.pages.pop(task_id, None)