diff --git a/skyvern/forge/sdk/db/client.py b/skyvern/forge/sdk/db/client.py index 19e2abde..50203542 100644 --- a/skyvern/forge/sdk/db/client.py +++ b/skyvern/forge/sdk/db/client.py @@ -1,4 +1,3 @@ -import asyncio import json from datetime import datetime, timedelta from typing import Any, List, Sequence @@ -3354,21 +3353,3 @@ class AgentDB: query = query.filter_by(organization_id=organization_id) task_run = (await session.scalars(query)).first() return Run.model_validate(task_run) if task_run else None - - async def wait_on_persistent_browser_address(self, session_id: str, organization_id: str) -> str: - async with asyncio.timeout(10 * 60): - while True: - persistent_browser_session = await self.get_persistent_browser_session(session_id, organization_id) - if persistent_browser_session is None: - raise Exception(f"Persistent browser session not found for {session_id}") - - LOG.info( - "Checking browser address", - session_id=session_id, - address=persistent_browser_session.browser_address, - ) - - if persistent_browser_session.browser_address: - return persistent_browser_session.browser_address - - await asyncio.sleep(2) diff --git a/skyvern/forge/sdk/db/polls.py b/skyvern/forge/sdk/db/polls.py index 02bc980b..addee714 100644 --- a/skyvern/forge/sdk/db/polls.py +++ b/skyvern/forge/sdk/db/polls.py @@ -3,13 +3,31 @@ import asyncio from structlog import get_logger from skyvern.forge.sdk.db.client import AgentDB +from skyvern.forge.sdk.schemas.persistent_browser_sessions import PersistentBrowserSession LOG = get_logger(__name__) -async def wait_on_persistent_browser_address(db: AgentDB, session_id: str, organization_id: str) -> str | None: +async def wait_on_persistent_browser_address( + db: AgentDB, + session_id: str, + organization_id: str, + timeout: int = 600, + poll_interval: int = 2, +) -> str | None: + persistent_browser_session = await await_browser_session(db, session_id, organization_id, timeout, poll_interval) + return persistent_browser_session.browser_address if persistent_browser_session else None + + +async def await_browser_session( + db: AgentDB, + session_id: str, + organization_id: str, + timeout: int = 600, + poll_interval: int = 2, +) -> PersistentBrowserSession | None: try: - async with asyncio.timeout(10 * 60): + async with asyncio.timeout(timeout): while True: persistent_browser_session = await db.get_persistent_browser_session(session_id, organization_id) if persistent_browser_session is None: @@ -22,9 +40,9 @@ async def wait_on_persistent_browser_address(db: AgentDB, session_id: str, organ ) if persistent_browser_session.browser_address: - return persistent_browser_session.browser_address + return persistent_browser_session - await asyncio.sleep(2) + await asyncio.sleep(poll_interval) except asyncio.TimeoutError: LOG.warning(f"Browser address not found for persistent browser session {session_id}") diff --git a/skyvern/forge/sdk/routes/streaming_verify.py b/skyvern/forge/sdk/routes/streaming_verify.py index 24d9ab61..f6e47943 100644 --- a/skyvern/forge/sdk/routes/streaming_verify.py +++ b/skyvern/forge/sdk/routes/streaming_verify.py @@ -52,11 +52,10 @@ async def verify_browser_session( ) try: - _, host, cdp_port = await app.PERSISTENT_SESSIONS_MANAGER.get_browser_address( + browser_address = await app.PERSISTENT_SESSIONS_MANAGER.get_browser_address( session_id=browser_session_id, organization_id=organization_id, ) - browser_address = f"{host}:{cdp_port}" except Exception as ex: LOG.info( "Browser session address not found for browser session.", @@ -211,11 +210,10 @@ async def verify_workflow_run( ) try: - _, host, cdp_port = await app.PERSISTENT_SESSIONS_MANAGER.get_browser_address( + browser_address = await app.PERSISTENT_SESSIONS_MANAGER.get_browser_address( session_id=browser_session.persistent_browser_session_id, organization_id=organization_id, ) - browser_address = f"{host}:{cdp_port}" except Exception as ex: LOG.info( "Browser session address not found for workflow run.", diff --git a/skyvern/webeye/persistent_sessions_manager.py b/skyvern/webeye/persistent_sessions_manager.py index faf5fa46..752fd832 100644 --- a/skyvern/webeye/persistent_sessions_manager.py +++ b/skyvern/webeye/persistent_sessions_manager.py @@ -5,6 +5,7 @@ from dataclasses import dataclass import structlog from playwright._impl._errors import TargetClosedError +from skyvern.exceptions import MissingBrowserAddressError 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 @@ -69,16 +70,13 @@ class PersistentSessionsManager: LOG.info("Browser session begin", browser_session_id=browser_session_id) - async def get_browser_address(self, session_id: str, organization_id: str) -> tuple[str, str, str]: + async def get_browser_address(self, session_id: str, organization_id: str) -> str: address = await wait_on_persistent_browser_address(self.database, session_id, organization_id) if address is None: - raise Exception(f"Browser address not found for persistent browser session {session_id}") + raise MissingBrowserAddressError(session_id) - protocol = "http" - host, cdp_port = address.split(":") - - return protocol, host, cdp_port + return address async def get_session_by_runnable_id( self, runnable_id: str, organization_id: str