support proxy select for pbs (#3461)

alembic/versions/2025_09_18_0916-4937dff73027_add_browser_session_proxy_schema.py

@@ -0,0 +1,31 @@
+"""add_browser_session_proxy_schema
+
+Revision ID: 4937dff73027
+Revises: 8f208daee7d1
+Create Date: 2025-09-18 09:16:54.130840+00:00
+
+"""
+
+from typing import Sequence, Union
+
+import sqlalchemy as sa
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "4937dff73027"
+down_revision: Union[str, None] = "8f208daee7d1"
+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("proxy_location", sa.String(), nullable=True))
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_column("persistent_browser_sessions", "proxy_location")
+    # ### end Alembic commands ###

skyvern/forge/sdk/db/client.py

@@ -3181,6 +3181,7 @@ class AgentDB:
         runnable_type: str | None = None,
         runnable_id: str | None = None,
         timeout_minutes: int | None = None,
+        proxy_location: ProxyLocation | None = ProxyLocation.RESIDENTIAL,
     ) -> PersistentBrowserSession:
         """Create a new persistent browser session."""
         try:
@@ -3190,6 +3191,7 @@ class AgentDB:
                     runnable_type=runnable_type,
                     runnable_id=runnable_id,
                     timeout_minutes=timeout_minutes,
+                    proxy_location=proxy_location,
                 )
                 session.add(browser_session)
                 await session.commit()

skyvern/forge/sdk/db/models.py

@@ -747,6 +747,7 @@ class PersistentBrowserSessionModel(Base):
     timeout_minutes = Column(Integer, nullable=True)
     ip_address = Column(String, nullable=True)
     ecs_task_arn = Column(String, nullable=True)
+    proxy_location = Column(String, 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/agent_protocol.py

@@ -82,6 +82,7 @@ from skyvern.schemas.runs import (
     CUA_ENGINES,
     BlockRunRequest,
     BlockRunResponse,
+    ProxyLocation,
     RunEngine,
     RunResponse,
     RunType,
@@ -2332,6 +2333,7 @@ async def new_debug_session(
     new_browser_session = await app.PERSISTENT_SESSIONS_MANAGER.create_session(
         organization_id=current_org.organization_id,
         timeout_minutes=settings.DEBUG_SESSION_TIMEOUT_MINUTES,
+        proxy_location=ProxyLocation.RESIDENTIAL_ISP,
     )
 
     debug_session = await app.DATABASE.create_debug_session(

skyvern/forge/sdk/routes/browser_sessions.py

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

skyvern/forge/sdk/schemas/persistent_browser_sessions.py

@@ -2,6 +2,8 @@ from datetime import datetime
 
 from pydantic import BaseModel, ConfigDict
 
+from skyvern.schemas.runs import ProxyLocation
+
 FINAL_STATUSES = ("completed", "failed")
 
 
@@ -20,6 +22,7 @@ class PersistentBrowserSession(BaseModel):
     ip_address: str | None = None
     status: str | None = None
     timeout_minutes: int | None = None
+    proxy_location: ProxyLocation | None = None
     started_at: datetime | None = None
     completed_at: datetime | None = None
     created_at: datetime

skyvern/schemas/browser_sessions.py

@@ -1,5 +1,8 @@
 from pydantic import BaseModel, Field
 
+from skyvern.schemas.docs.doc_strings import PROXY_LOCATION_DOC_STRING
+from skyvern.schemas.runs import ProxyLocation
+
 MIN_TIMEOUT = 5
 MAX_TIMEOUT = 60 * 4  # 4 hours
 DEFAULT_TIMEOUT = 60
@@ -12,3 +15,7 @@ class CreateBrowserSessionRequest(BaseModel):
         ge=MIN_TIMEOUT,
         le=MAX_TIMEOUT,
     )
+    proxy_location: ProxyLocation | None = Field(
+        default=None,
+        description=PROXY_LOCATION_DOC_STRING,
+    )

skyvern/webeye/persistent_sessions_manager.py

@@ -12,6 +12,7 @@ from skyvern.exceptions import BrowserSessionNotRenewable, MissingBrowserAddress
 from skyvern.forge.sdk.db.client import AgentDB
 from skyvern.forge.sdk.db.polls import wait_on_persistent_browser_address
 from skyvern.forge.sdk.schemas.persistent_browser_sessions import PersistentBrowserSession, is_final_status
+from skyvern.schemas.runs import ProxyLocation
 from skyvern.webeye.browser_factory import BrowserState
 
 LOG = structlog.get_logger()
@@ -237,6 +238,7 @@ class PersistentSessionsManager:
         runnable_id: str | None = None,
         runnable_type: str | None = None,
         timeout_minutes: int | None = None,
+        proxy_location: ProxyLocation | None = ProxyLocation.RESIDENTIAL,
     ) -> PersistentBrowserSession:
         """Create a new browser session for an organization and return its ID with the browser state."""
 
@@ -250,6 +252,7 @@ class PersistentSessionsManager:
             runnable_type=runnable_type,
             runnable_id=runnable_id,
             timeout_minutes=timeout_minutes,
+            proxy_location=proxy_location,
         )
 
         return browser_session_db