pass organization id to launch browser (#1000)

This commit is contained in:
LawyZheng
2024-10-18 15:02:11 +08:00
committed by GitHub
parent f69016088b
commit 533c814558
2 changed files with 37 additions and 7 deletions

View File

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

View File

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