From d3dbdee800488864398d044e73eb696c0f956fd0 Mon Sep 17 00:00:00 2001 From: Shuchang Zheng Date: Thu, 20 Mar 2025 14:26:14 -0700 Subject: [PATCH] Add browser session id permission checking for task v1, v2, and workflow runs (#1980) --- .../core/permissions/permission_checkers.py | 4 ++-- skyvern/forge/sdk/db/client.py | 19 +++++++++++-------- skyvern/forge/sdk/routes/agent_protocol.py | 4 +++- skyvern/webeye/browser_manager.py | 8 ++++++-- skyvern/webeye/persistent_sessions_manager.py | 2 +- 5 files changed, 23 insertions(+), 14 deletions(-) diff --git a/skyvern/forge/sdk/core/permissions/permission_checkers.py b/skyvern/forge/sdk/core/permissions/permission_checkers.py index 01b9a748..eebdd54c 100644 --- a/skyvern/forge/sdk/core/permissions/permission_checkers.py +++ b/skyvern/forge/sdk/core/permissions/permission_checkers.py @@ -5,10 +5,10 @@ from skyvern.forge.sdk.schemas.organizations import Organization class PermissionChecker(abc.ABC): @abc.abstractmethod - async def check(self, organization: Organization) -> None: + async def check(self, organization: Organization, browser_session_id: str | None = None) -> None: pass class NoopPermissionChecker(PermissionChecker): - async def check(self, organization: Organization) -> None: + async def check(self, organization: Organization, browser_session_id: str | None = None) -> None: return diff --git a/skyvern/forge/sdk/db/client.py b/skyvern/forge/sdk/db/client.py index e4747ead..1df80eb5 100644 --- a/skyvern/forge/sdk/db/client.py +++ b/skyvern/forge/sdk/db/client.py @@ -2585,17 +2585,20 @@ class AgentDB: LOG.error("UnexpectedError", exc_info=True) raise - async def get_persistent_browser_session_by_id(self, session_id: str) -> Optional[PersistentBrowserSession]: + async def get_persistent_browser_session_by_id( + self, session_id: str, organization_id: str | None = None + ) -> Optional[PersistentBrowserSession]: """Get a specific persistent browser session.""" try: async with self.Session() as session: - persistent_browser_session = ( - await session.scalars( - select(PersistentBrowserSessionModel) - .filter_by(persistent_browser_session_id=session_id) - .filter_by(deleted_at=None) - ) - ).first() + query = ( + select(PersistentBrowserSessionModel) + .filter_by(persistent_browser_session_id=session_id) + .filter_by(deleted_at=None) + ) + if organization_id: + query = query.filter_by(organization_id=organization_id) + persistent_browser_session = (await session.scalars(query)).first() if persistent_browser_session: return PersistentBrowserSession.model_validate(persistent_browser_session) raise NotFoundError(f"PersistentBrowserSession {session_id} not found") diff --git a/skyvern/forge/sdk/routes/agent_protocol.py b/skyvern/forge/sdk/routes/agent_protocol.py index 415d41a9..3f8c6307 100644 --- a/skyvern/forge/sdk/routes/agent_protocol.py +++ b/skyvern/forge/sdk/routes/agent_protocol.py @@ -168,7 +168,7 @@ async def run_task( x_max_steps_override: Annotated[int | None, Header()] = None, ) -> CreateTaskResponse: analytics.capture("skyvern-oss-agent-task-create", data={"url": task.url}) - await PermissionCheckerFactory.get_instance().check(current_org) + await PermissionCheckerFactory.get_instance().check(current_org, browser_session_id=task.browser_session_id) created_task = await app.agent.create_task(task, current_org.organization_id) url_hash = generate_url_hash(task.url) @@ -574,6 +574,7 @@ async def run_workflow( analytics.capture("skyvern-oss-agent-workflow-execute") context = skyvern_context.ensure_context() request_id = context.request_id + await PermissionCheckerFactory.get_instance().check(current_org) if template: if workflow_id not in await app.STORAGE.retrieve_global_workflows(): @@ -1121,6 +1122,7 @@ async def run_task_v2( max_iterations_override=x_max_iterations_override, max_steps_override=x_max_steps_override, ) + await PermissionCheckerFactory.get_instance().check(organization, browser_session_id=data.browser_session_id) try: task_v2 = await task_v2_service.initialize_task_v2( diff --git a/skyvern/webeye/browser_manager.py b/skyvern/webeye/browser_manager.py index 3389d10a..0a733010 100644 --- a/skyvern/webeye/browser_manager.py +++ b/skyvern/webeye/browser_manager.py @@ -81,7 +81,9 @@ class BrowserManager: "Getting browser state for task from persistent sessions manager", browser_session_id=browser_session_id, ) - browser_state = await app.PERSISTENT_SESSIONS_MANAGER.get_browser_state(browser_session_id) + browser_state = await app.PERSISTENT_SESSIONS_MANAGER.get_browser_state( + browser_session_id, organization_id=task.organization_id + ) if browser_state is None: LOG.warning( "Browser state not found in persistent sessions manager", @@ -148,7 +150,9 @@ class BrowserManager: "Getting browser state for workflow run from persistent sessions manager", browser_session_id=browser_session_id, ) - browser_state = await app.PERSISTENT_SESSIONS_MANAGER.get_browser_state(browser_session_id) + browser_state = await app.PERSISTENT_SESSIONS_MANAGER.get_browser_state( + browser_session_id, organization_id=workflow_run.organization_id + ) if browser_state is None: LOG.warning( "Browser state not found in persistent sessions manager", browser_session_id=browser_session_id diff --git a/skyvern/webeye/persistent_sessions_manager.py b/skyvern/webeye/persistent_sessions_manager.py index bf840a9e..2ffab70c 100644 --- a/skyvern/webeye/persistent_sessions_manager.py +++ b/skyvern/webeye/persistent_sessions_manager.py @@ -35,7 +35,7 @@ class PersistentSessionsManager: """Get all active sessions for an organization.""" return await self.database.get_active_persistent_browser_sessions(organization_id) - async def get_browser_state(self, session_id: str) -> BrowserState | None: + async def get_browser_state(self, session_id: str, organization_id: str | None = None) -> BrowserState | None: """Get a specific browser session's state by session ID.""" browser_session = self._browser_sessions.get(session_id) return browser_session.browser_state if browser_session else None