support extension choice in pbs (#4364)

alembic/versions/2025_12_24_0456-e393f33ec711_add_extension_for_browser_session.py

@@ -0,0 +1,31 @@
+"""add extension for browser session
+
+Revision ID: e393f33ec711
+Revises: b4738bd17198
+Create Date: 2025-12-24 04:56:38.500907+00:00
+
+"""
+
+from typing import Sequence, Union
+
+import sqlalchemy as sa
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "e393f33ec711"
+down_revision: Union[str, None] = "b4738bd17198"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column("persistent_browser_sessions", sa.Column("extensions", sa.JSON(), nullable=True))
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_column("persistent_browser_sessions", "extensions")
+    # ### end Alembic commands ###

skyvern/forge/sdk/db/agent_db.py

@@ -94,7 +94,7 @@ from skyvern.forge.sdk.schemas.organizations import (
     Organization,
     OrganizationAuthToken,
 )
-from skyvern.forge.sdk.schemas.persistent_browser_sessions import PersistentBrowserSession
+from skyvern.forge.sdk.schemas.persistent_browser_sessions import Extensions, PersistentBrowserSession
 from skyvern.forge.sdk.schemas.runs import Run
 from skyvern.forge.sdk.schemas.task_generations import TaskGeneration
 from skyvern.forge.sdk.schemas.task_v2 import TaskV2, TaskV2Status, Thought, ThoughtType
@@ -4540,8 +4540,10 @@ class AgentDB(BaseAlchemyDB):
         runnable_id: str | None = None,
         timeout_minutes: int | None = None,
         proxy_location: ProxyLocationInput = ProxyLocation.RESIDENTIAL,
+        extensions: list[Extensions] | None = None,
     ) -> PersistentBrowserSession:
         """Create a new persistent browser session."""
+        extensions_str: list[str] | None = [extension.value for extension in extensions] if extensions else None
         try:
             async with self.Session() as session:
                 browser_session = PersistentBrowserSessionModel(
@@ -4550,6 +4552,7 @@ class AgentDB(BaseAlchemyDB):
                     runnable_id=runnable_id,
                     timeout_minutes=timeout_minutes,
                     proxy_location=_serialize_proxy_location(proxy_location),
+                    extensions=extensions_str,
                 )
                 session.add(browser_session)
                 await session.commit()

skyvern/forge/sdk/db/models.py

@@ -851,6 +851,7 @@ class PersistentBrowserSessionModel(Base):
     ip_address = Column(String, nullable=True)
     ecs_task_arn = Column(String, nullable=True)
     proxy_location = Column(String, nullable=True)
+    extensions = Column(JSON, nullable=True)
     started_at = Column(DateTime, nullable=True)
     completed_at = Column(DateTime, nullable=True)
     created_at = Column(DateTime, default=datetime.datetime.utcnow, nullable=False, index=True)

skyvern/forge/sdk/routes/browser_sessions.py

@@ -93,6 +93,7 @@ async def create_browser_session(
         organization_id=current_org.organization_id,
         timeout_minutes=browser_session_request.timeout,
         proxy_location=browser_session_request.proxy_location,
+        extensions=browser_session_request.extensions,
     )
     return await BrowserSessionResponse.from_browser_session(browser_session)
 

skyvern/forge/sdk/schemas/persistent_browser_sessions.py

@@ -26,6 +26,11 @@ def is_final_status(status: str | None) -> bool:
     return status in FINAL_STATUSES
 
 
+class Extensions(StrEnum):
+    AdBlocker = "ad-blocker"
+    CaptchaSolver = "captcha-solver"
+
+
 class PersistentBrowserSession(BaseModel):
     model_config = ConfigDict(from_attributes=True)
 
@@ -43,6 +48,7 @@ class PersistentBrowserSession(BaseModel):
     created_at: datetime
     modified_at: datetime
     deleted_at: datetime | None = None
+    extensions: list[Extensions] | None = None
 
 
 class AddressablePersistentBrowserSession(PersistentBrowserSession):

skyvern/schemas/browser_sessions.py

@@ -2,6 +2,7 @@ from pydantic import BaseModel, Field
 
 from skyvern.client.types.workflow_definition_yaml_blocks_item import WorkflowDefinitionYamlBlocksItem
 from skyvern.client.types.workflow_definition_yaml_parameters_item import WorkflowDefinitionYamlParametersItem_Workflow
+from skyvern.forge.sdk.schemas.persistent_browser_sessions import Extensions
 from skyvern.schemas.docs.doc_strings import PROXY_LOCATION_DOC_STRING
 from skyvern.schemas.runs import ProxyLocation
 
@@ -22,6 +23,11 @@ class CreateBrowserSessionRequest(BaseModel):
         description=PROXY_LOCATION_DOC_STRING,
     )
 
+    extensions: list[Extensions] | None = Field(
+        default=None,
+        description="A list of extensions to install in the browser session.",
+    )
+
 
 class ProcessBrowserSessionRecordingRequest(BaseModel):
     compressed_chunks: list[str] = Field(

skyvern/webeye/persistent_sessions_manager.py

@@ -13,6 +13,7 @@ from skyvern.forge import app
 from skyvern.forge.sdk.db.agent_db import AgentDB
 from skyvern.forge.sdk.db.polls import wait_on_persistent_browser_address
 from skyvern.forge.sdk.schemas.persistent_browser_sessions import (
+    Extensions,
     PersistentBrowserSession,
     PersistentBrowserSessionStatus,
     is_final_status,
@@ -256,6 +257,7 @@ class PersistentSessionsManager:
         runnable_type: str | None = None,
         timeout_minutes: int | None = None,
         proxy_location: ProxyLocationInput = ProxyLocation.RESIDENTIAL,
+        extensions: list[Extensions] | None = None,
     ) -> PersistentBrowserSession:
         """Create a new browser session for an organization and return its ID with the browser state."""
 
@@ -270,6 +272,7 @@ class PersistentSessionsManager:
             runnable_id=runnable_id,
             timeout_minutes=timeout_minutes,
             proxy_location=proxy_location,
+            extensions=extensions,
         )
 
         return browser_session_db

skyvern/webeye/schemas.py

@@ -10,7 +10,7 @@ from skyvern.config import settings
 from skyvern.constants import GET_DOWNLOADED_FILES_TIMEOUT
 from skyvern.forge.sdk.artifact.storage.base import BaseStorage
 from skyvern.forge.sdk.schemas.files import FileInfo
-from skyvern.forge.sdk.schemas.persistent_browser_sessions import PersistentBrowserSession
+from skyvern.forge.sdk.schemas.persistent_browser_sessions import Extensions, PersistentBrowserSession
 
 LOG = structlog.get_logger()
 
@@ -46,6 +46,10 @@ class BrowserSessionResponse(BaseModel):
         description="Url for the browser session page",
         examples=["https://app.skyvern.com/browser-session/pbs_123456"],
     )
+    extensions: list[Extensions] | None = Field(
+        None,
+        description="A list of extensions installed in the browser session.",
+    )
     vnc_streaming_supported: bool = Field(False, description="Whether the browser session supports VNC streaming")
     download_path: str | None = Field(None, description="The path where the browser session downloads files")
     downloaded_files: list[FileInfo] | None = Field(
@@ -126,4 +130,5 @@ class BrowserSessionResponse(BaseModel):
             download_path=download_path,
             downloaded_files=downloaded_files,
             recordings=recordings,
+            extensions=browser_session.extensions,
         )