From 533c81455882d42c3fee10d6eb925bbf80e717f9 Mon Sep 17 00:00:00 2001
From: LawyZheng <lawyzheng1106@gmail.com>
Date: Fri, 18 Oct 2024 15:02:11 +0800
Subject: [PATCH] pass organization id to launch browser (#1000)

---
 skyvern/webeye/browser_factory.py | 21 ++++++++++++++++++---
 skyvern/webeye/browser_manager.py | 23 +++++++++++++++++++----
 2 files changed, 37 insertions(+), 7 deletions(-)

diff --git a/skyvern/webeye/browser_factory.py b/skyvern/webeye/browser_factory.py
index 19a3cc40..faec9357 100644
--- a/skyvern/webeye/browser_factory.py
+++ b/skyvern/webeye/browser_factory.py
@@ -208,6 +208,7 @@ class BrowserState:
         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 @@ class BrowserState:
                 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,6 +314,7 @@ class BrowserState:
         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:
@@ -319,7 +322,11 @@ class BrowserState:
 
         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)
@@ -327,14 +334,22 @@ class BrowserState:
                 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
index 5ce4d460..055eb4e4 100644
--- a/skyvern/webeye/browser_manager.py
+++ b/skyvern/webeye/browser_manager.py
@@ -30,6 +30,7 @@ class BrowserManager:
         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 @@ class BrowserManager:
             url=url,
             task_id=task_id,
             workflow_run_id=workflow_run_id,
+            organization_id=organization_id,
         )
         return BrowserState(
             pw=pw,
@@ -64,11 +66,18 @@ class BrowserManager:
             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 @@ class BrowserManager:
             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