browser session streaming fix and clean up wait_on_persistent_browser_address (#2983)

This commit is contained in:
Shuchang Zheng
2025-07-18 00:16:07 -07:00
committed by GitHub
parent 8263c97e16
commit 764ac1928c
4 changed files with 28 additions and 33 deletions

View File

@@ -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)

View File

@@ -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}")

View File

@@ -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.",