From dde3dd55f62dd2924e2ca0c30a81988628cefec4 Mon Sep 17 00:00:00 2001 From: Maksim Ivanov Date: Wed, 18 Dec 2024 07:29:30 +0100 Subject: [PATCH] Store persistent browser sessions data in the db --- skyvern/forge/sdk/db/client.py | 14 ++++---- skyvern/forge/sdk/routes/agent_protocol.py | 6 ++-- skyvern/webeye/models.py | 10 ++++-- skyvern/webeye/persistent_sessions_manager.py | 35 +++++++++++-------- 4 files changed, 39 insertions(+), 26 deletions(-) diff --git a/skyvern/forge/sdk/db/client.py b/skyvern/forge/sdk/db/client.py index cb0ddf166..aa861871d 100644 --- a/skyvern/forge/sdk/db/client.py +++ b/skyvern/forge/sdk/db/client.py @@ -1879,7 +1879,7 @@ async def update_observer_cruise( return ObserverCruise.model_validate(observer_cruise) raise NotFoundError(f"ObserverCruise {observer_cruise_id} not found") - async def get_active_persistent_browser_session_ids(self, organization_id: str) -> List[str]: + async def get_active_persistent_browser_sessions(self, organization_id: str) -> List[str]: """Get all active persistent browser session IDs for an organization.""" async with self.Session() as session: result = await session.execute( @@ -1888,7 +1888,7 @@ async def get_active_persistent_browser_session_ids(self, organization_id: str) PersistentBrowserSessionModel.deleted_at.is_(None), ) ) - return [row[0] for row in result.all()] + return result.scalars().all() async def get_persistent_browser_session(self, session_id: str, organization_id: str) -> Optional[PersistentBrowserSessionModel]: """Get a specific persistent browser session.""" @@ -1902,15 +1902,17 @@ async def get_persistent_browser_session(self, session_id: str, organization_id: return result.scalar_one_or_none() async def create_persistent_browser_session( - self, session_id: str, organization_id: str + self, + organization_id: str, + runnable_type: str, + runnable_id: str, ) -> PersistentBrowserSessionModel: """Create a new persistent browser session.""" async with self.Session() as session: db_session = PersistentBrowserSessionModel( - persistent_browser_session_id=session_id, organization_id=organization_id, - runnable_type="browser_session", - runnable_id=session_id, + runnable_type=runnable_type, + runnable_id=runnable_id, ) session.add(db_session) await session.commit() diff --git a/skyvern/forge/sdk/routes/agent_protocol.py b/skyvern/forge/sdk/routes/agent_protocol.py index e07157c9e..0701b4b84 100644 --- a/skyvern/forge/sdk/routes/agent_protocol.py +++ b/skyvern/forge/sdk/routes/agent_protocol.py @@ -990,13 +990,13 @@ async def get_browser_sessions( ) -> list[BrowserSessionResponse]: """Get all active browser sessions for the organization""" analytics.capture("skyvern-oss-agent-browser-sessions-get") - session_ids = await app.PERSISTENT_SESSIONS_MANAGER.get_active_session_ids(current_org.organization_id) + browser_sessions = await app.PERSISTENT_SESSIONS_MANAGER.get_active_sessions(current_org.organization_id) return [ await app.PERSISTENT_SESSIONS_MANAGER.build_browser_session_response( organization_id=current_org.organization_id, - session_id=session_id, + session_id=browser_session.persistent_browser_session_id, ) - for session_id in session_ids + for browser_session in browser_sessions ] @base_router.post( diff --git a/skyvern/webeye/models.py b/skyvern/webeye/models.py index c8c0fde18..62200058d 100644 --- a/skyvern/webeye/models.py +++ b/skyvern/webeye/models.py @@ -1,8 +1,12 @@ -from pydantic import BaseModel - -from skyvern.webeye.browser_factory import BrowserState +from datetime import datetime +from pydantic import BaseModel class BrowserSessionResponse(BaseModel): session_id: str organization_id: str + runnable_type: str + runnable_id: str + created_at: datetime + modified_at: datetime + deleted_at: datetime | None diff --git a/skyvern/webeye/persistent_sessions_manager.py b/skyvern/webeye/persistent_sessions_manager.py index 8419e5926..ac59c1d73 100644 --- a/skyvern/webeye/persistent_sessions_manager.py +++ b/skyvern/webeye/persistent_sessions_manager.py @@ -18,7 +18,6 @@ class PersistentSessionsManager: instance = None - # Store BrowserState objects in memory since they can't be serialized to DB _browser_states: Dict[str, BrowserState] = dict() def __init__(self, database: AgentDB): @@ -30,12 +29,12 @@ def __new__(cls, database: AgentDB) -> PersistentSessionsManager: cls.instance.database = database return cls.instance - async def get_active_session_ids(self, organization_id: str) -> List[str]: + async def get_active_sessions(self, organization_id: str) -> List[str]: """Get all active session IDs for an organization.""" - return await self.database.get_active_persistent_browser_session_ids(organization_id) + return await self.database.get_active_persistent_browser_sessions(organization_id) - def get_session(self, organization_id: str, session_id: str) -> Optional[BrowserState]: - """Get a specific browser session by organization ID and session ID.""" + def get_session(self, session_id: str) -> Optional[BrowserState]: + """Get a specific browser session by session ID.""" return self._browser_states.get(session_id) async def create_session( @@ -43,18 +42,23 @@ async def create_session( organization_id: str, proxy_location: ProxyLocation | None = None, url: str | None = None, + runnable_id: str | None = None, + runnable_type: str | None = None, ) -> Tuple[str, BrowserState]: """Create a new browser session for an organization and return its ID with the browser state.""" - session_id = generate_persistent_browser_session_id() LOG.info( "Creating new browser session", organization_id=organization_id, - session_id=session_id, ) - # Create database record - await self.database.create_persistent_browser_session(session_id, organization_id) + browser_session = await self.database.create_persistent_browser_session( + organization_id=organization_id, + runnable_type=runnable_type, + runnable_id=runnable_id, + ) + + session_id = browser_session.persistent_browser_session_id pw = await async_playwright().start() browser_context, browser_artifacts, browser_cleanup = await BrowserContextFactory.create_browser_context( @@ -79,7 +83,6 @@ async def on_context_close(): self._browser_states[session_id] = browser_state - # Create initial page if URL is provided if url: await browser_state.get_or_create_page( url=url, @@ -99,7 +102,6 @@ async def close_session(self, organization_id: str, session_id: str) -> None: session_id=session_id, ) await browser_state.close() - self._browser_states.pop(session_id, None) # Mark as deleted in database await self.database.mark_persistent_browser_session_deleted(session_id, organization_id) @@ -110,10 +112,15 @@ async def close_all_sessions(self, organization_id: str) -> None: for session_id in session_ids: await self.close_session(organization_id, session_id) - async def build_browser_session_response(self, organization_id: str, session_id: str) -> BrowserSessionResponse: + async def build_browser_session_response(self, browser_session: PersistentBrowserSessionModel) -> BrowserSessionResponse: return BrowserSessionResponse( - session_id=session_id, - organization_id=organization_id, + session_id=browser_session.persistent_browser_session_id, + organization_id=browser_session.organization_id, + runnable_type=browser_session.runnable_type, + runnable_id=browser_session.runnable_id, + created_at=browser_session.created_at, + modified_at=browser_session.modified_at, + deleted_at=browser_session.deleted_at, ) @classmethod