pass organization id to launch browser

Skyvern-AI · Oct 18, 2024 · 75060df · 75060df
1 parent f690160
commit 75060df
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 7 deletions.
diff --git a/skyvern/webeye/browser_factory.py b/skyvern/webeye/browser_factory.py
@@ -208,6 +208,7 @@ async def check_and_fix_state(
         proxy_location: ProxyLocation | None = None,
         task_id: str | None = None,
         workflow_run_id: str | None = None,
+        organization_id: str | None = None,
     ) -> None:
         if self.pw is None:
             LOG.info("Starting playwright")
@@ -225,6 +226,7 @@ async def check_and_fix_state(
                 proxy_location=proxy_location,
                 task_id=task_id,
                 workflow_run_id=workflow_run_id,
+                organization_id=organization_id,
             )
             self.browser_context = browser_context
             self.browser_artifacts = browser_artifacts
@@ -312,29 +314,42 @@ async def get_or_create_page(
         proxy_location: ProxyLocation | None = None,
         task_id: str | None = None,
         workflow_run_id: str | None = None,
+        organization_id: str | None = None,
     ) -> Page:
         page = await self.get_working_page()
         if page is not None:
             return page
 
         try:
             await self.check_and_fix_state(
-                url=url, proxy_location=proxy_location, task_id=task_id, workflow_run_id=workflow_run_id
+                url=url,
+                proxy_location=proxy_location,
+                task_id=task_id,
+                workflow_run_id=workflow_run_id,
+                organization_id=organization_id,
             )
         except Exception as e:
             error_message = str(e)
             if "net::ERR" not in error_message:
                 raise e
             await self.close_current_open_page()
             await self.check_and_fix_state(
-                url=url, proxy_location=proxy_location, task_id=task_id, workflow_run_id=workflow_run_id
+                url=url,
+                proxy_location=proxy_location,
+                task_id=task_id,
+                workflow_run_id=workflow_run_id,
+                organization_id=organization_id,
             )
         await self.__assert_page()
 
         if not await BrowserContextFactory.validate_browser_context(await self.get_working_page()):
             await self.close_current_open_page()
             await self.check_and_fix_state(
-                url=url, proxy_location=proxy_location, task_id=task_id, workflow_run_id=workflow_run_id
+                url=url,
+                proxy_location=proxy_location,
+                task_id=task_id,
+                workflow_run_id=workflow_run_id,
+                organization_id=organization_id,
             )
             await self.__assert_page()
 

diff --git a/skyvern/webeye/browser_manager.py b/skyvern/webeye/browser_manager.py
@@ -30,6 +30,7 @@ async def _create_browser_state(
         url: str | None = None,
         task_id: str | None = None,
         workflow_run_id: str | None = None,
+        organization_id: str | None = None,
     ) -> BrowserState:
         pw = await async_playwright().start()
         (
@@ -42,6 +43,7 @@ async def _create_browser_state(
             url=url,
             task_id=task_id,
             workflow_run_id=workflow_run_id,
+            organization_id=organization_id,
         )
         return BrowserState(
             pw=pw,
@@ -64,11 +66,18 @@ async def get_or_create_for_task(self, task: Task) -> BrowserState:
             return self.pages[task.task_id]
 
         LOG.info("Creating browser state for task", task_id=task.task_id)
-        browser_state = await self._create_browser_state(task.proxy_location, task.url, task.task_id)
+        browser_state = await self._create_browser_state(
+            proxy_location=task.proxy_location,
+            url=task.url,
+            task_id=task.task_id,
+            organization_id=task.organization_id,
+        )
 
         # The URL here is only used when creating a new page, and not when using an existing page.
         # This will make sure browser_state.page is not None.
-        await browser_state.get_or_create_page(url=task.url, proxy_location=task.proxy_location, task_id=task.task_id)
+        await browser_state.get_or_create_page(
+            url=task.url, proxy_location=task.proxy_location, task_id=task.task_id, organization_id=task.organization_id
+        )
 
         self.pages[task.task_id] = browser_state
         if task.workflow_run_id:
@@ -83,13 +92,19 @@ async def get_or_create_for_workflow_run(self, workflow_run: WorkflowRun, url: s
             workflow_run_id=workflow_run.workflow_run_id,
         )
         browser_state = await self._create_browser_state(
-            workflow_run.proxy_location, url=url, workflow_run_id=workflow_run.workflow_run_id
+            workflow_run.proxy_location,
+            url=url,
+            workflow_run_id=workflow_run.workflow_run_id,
+            organization_id=workflow_run.organization_id,
         )
 
         # The URL here is only used when creating a new page, and not when using an existing page.
         # This will make sure browser_state.page is not None.
         await browser_state.get_or_create_page(
-            url=url, proxy_location=workflow_run.proxy_location, workflow_run_id=workflow_run.workflow_run_id
+            url=url,
+            proxy_location=workflow_run.proxy_location,
+            workflow_run_id=workflow_run.workflow_run_id,
+            organization_id=workflow_run.organization_id,
         )
 
         self.pages[workflow_run.workflow_run_id] = browser_state