From 3e12133af96c5c1b599104d9e4b3c738f03df49f Mon Sep 17 00:00:00 2001 From: LawyZheng Date: Thu, 18 Sep 2025 17:21:46 +0800 Subject: [PATCH] support proxy select for pbs (#3461) --- ...f73027_add_browser_session_proxy_schema.py | 31 +++++++++++++++++++ skyvern/forge/sdk/db/client.py | 2 ++ skyvern/forge/sdk/db/models.py | 1 + skyvern/forge/sdk/routes/agent_protocol.py | 2 ++ skyvern/forge/sdk/routes/browser_sessions.py | 1 + .../schemas/persistent_browser_sessions.py | 3 ++ skyvern/schemas/browser_sessions.py | 7 +++++ skyvern/webeye/persistent_sessions_manager.py | 3 ++ 8 files changed, 50 insertions(+) create mode 100644 alembic/versions/2025_09_18_0916-4937dff73027_add_browser_session_proxy_schema.py diff --git a/alembic/versions/2025_09_18_0916-4937dff73027_add_browser_session_proxy_schema.py b/alembic/versions/2025_09_18_0916-4937dff73027_add_browser_session_proxy_schema.py new file mode 100644 index 00000000..3eebb992 --- /dev/null +++ b/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 ### diff --git a/skyvern/forge/sdk/db/client.py b/skyvern/forge/sdk/db/client.py index c378fa3b..619b675c 100644 --- a/skyvern/forge/sdk/db/client.py +++ b/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() diff --git a/skyvern/forge/sdk/db/models.py b/skyvern/forge/sdk/db/models.py index 232f01ea..3e607b01 100644 --- a/skyvern/forge/sdk/db/models.py +++ b/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) diff --git a/skyvern/forge/sdk/routes/agent_protocol.py b/skyvern/forge/sdk/routes/agent_protocol.py index b1e348e9..19cf3d39 100644 --- a/skyvern/forge/sdk/routes/agent_protocol.py +++ b/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( diff --git a/skyvern/forge/sdk/routes/browser_sessions.py b/skyvern/forge/sdk/routes/browser_sessions.py index d975d6d1..076c1343 100644 --- a/skyvern/forge/sdk/routes/browser_sessions.py +++ b/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) diff --git a/skyvern/forge/sdk/schemas/persistent_browser_sessions.py b/skyvern/forge/sdk/schemas/persistent_browser_sessions.py index 4b5b4f88..43c1f7c8 100644 --- a/skyvern/forge/sdk/schemas/persistent_browser_sessions.py +++ b/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 diff --git a/skyvern/schemas/browser_sessions.py b/skyvern/schemas/browser_sessions.py index 9139f576..02040c07 100644 --- a/skyvern/schemas/browser_sessions.py +++ b/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, + ) diff --git a/skyvern/webeye/persistent_sessions_manager.py b/skyvern/webeye/persistent_sessions_manager.py index a0ad09ef..0fabe50b 100644 --- a/skyvern/webeye/persistent_sessions_manager.py +++ b/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