Add persist_browser_session flag to workflows (#777)

skyvern/forge/sdk/artifact/storage/base.py

@@ -59,3 +59,11 @@ class BaseStorage(ABC):
     @abstractmethod
     async def get_streaming_file(self, organization_id: str, file_name: str, use_default: bool = True) -> bytes | None:
         pass
+
+    @abstractmethod
+    async def store_browser_session(self, organization_id: str, workflow_permanent_id: str, directory: str) -> None:
+        pass
+
+    @abstractmethod
+    async def retrieve_browser_session(self, organization_id: str, workflow_permanent_id: str) -> str | None:
+        pass

skyvern/forge/sdk/artifact/storage/local.py

@@ -1,3 +1,5 @@
+import os
+import shutil
 from datetime import datetime
 from pathlib import Path
 from urllib.parse import unquote, urlparse
@@ -85,6 +87,38 @@ class LocalStorage(BaseStorage):
             )
             return None
 
+    async def store_browser_session(self, organization_id: str, workflow_permanent_id: str, directory: str) -> None:
+        stored_folder_path = (
+            Path(SettingsManager.get_settings().BROWSER_SESSION_BASE_PATH) / organization_id / workflow_permanent_id
+        )
+        if directory == str(stored_folder_path):
+            return
+        self._create_directories_if_not_exists(stored_folder_path)
+        LOG.info(
+            "Storing browser session locally",
+            organization_id=organization_id,
+            workflow_permanent_id=workflow_permanent_id,
+            directory=directory,
+            browser_session_path=stored_folder_path,
+        )
+
+        # Copy all files from the directory to the stored folder
+        for root, _, files in os.walk(directory):
+            for file in files:
+                source_file_path = Path(root) / file
+                relative_path = source_file_path.relative_to(directory)
+                target_file_path = stored_folder_path / relative_path
+                self._create_directories_if_not_exists(target_file_path)
+                shutil.copy2(source_file_path, target_file_path)
+
+    async def retrieve_browser_session(self, organization_id: str, workflow_permanent_id: str) -> str | None:
+        stored_folder_path = (
+            Path(SettingsManager.get_settings().BROWSER_SESSION_BASE_PATH) / organization_id / workflow_permanent_id
+        )
+        if not stored_folder_path.exists():
+            return None
+        return str(stored_folder_path)
+
     @staticmethod
     def _parse_uri_to_path(uri: str) -> str:
         parsed_uri = urlparse(uri)

skyvern/forge/sdk/artifact/storage/s3.py

@@ -1,7 +1,10 @@
+import shutil
+import tempfile
 from datetime import datetime
 
 from skyvern.config import settings
 from skyvern.forge.sdk.api.aws import AsyncAWSClient
+from skyvern.forge.sdk.api.files import unzip_files
 from skyvern.forge.sdk.artifact.models import Artifact, ArtifactType
 from skyvern.forge.sdk.artifact.storage.base import FILE_EXTENTSION_MAP, BaseStorage
 from skyvern.forge.sdk.models import Step
@@ -40,3 +43,24 @@ class S3Storage(BaseStorage):
     async def get_streaming_file(self, organization_id: str, file_name: str, use_default: bool = True) -> bytes | None:
         path = f"s3://{settings.AWS_S3_BUCKET_SCREENSHOTS}/{settings.ENV}/{organization_id}/{file_name}"
         return await self.async_client.download_file(path, log_exception=False)
+
+    async def store_browser_session(self, organization_id: str, workflow_permanent_id: str, directory: str) -> None:
+        # Zip the directory to a temp file
+        temp_zip_file = tempfile.NamedTemporaryFile()
+        zip_file_path = shutil.make_archive(temp_zip_file.name, "zip", directory)
+        browser_session_uri = f"s3://{settings.AWS_S3_BUCKET_BROWSER_SESSIONS}/{settings.ENV}/{organization_id}/{workflow_permanent_id}.zip"
+        await self.async_client.upload_file_from_path(browser_session_uri, zip_file_path)
+
+    async def retrieve_browser_session(self, organization_id: str, workflow_permanent_id: str) -> str | None:
+        browser_session_uri = f"s3://{settings.AWS_S3_BUCKET_BROWSER_SESSIONS}/{settings.ENV}/{organization_id}/{workflow_permanent_id}.zip"
+        downloaded_zip_bytes = await self.async_client.download_file(browser_session_uri, log_exception=True)
+        if not downloaded_zip_bytes:
+            return None
+        temp_zip_file = tempfile.NamedTemporaryFile(delete=False)
+        temp_zip_file.write(downloaded_zip_bytes)
+        temp_zip_file_path = temp_zip_file.name
+
+        temp_dir = tempfile.mkdtemp(prefix="skyvern_browser_session_")
+        unzip_files(temp_zip_file_path, temp_dir)
+        temp_zip_file.close()
+        return temp_dir