user defined browser header (#2752)
Co-authored-by: lawyzheng <lawyzheng1106@gmail.com>
This commit is contained in:
@@ -170,7 +170,11 @@ class BrowserContextFactory:
|
||||
f.write(preference_file_content)
|
||||
|
||||
@staticmethod
|
||||
def build_browser_args(proxy_location: ProxyLocation | None = None, cdp_port: int | None = None) -> dict[str, Any]:
|
||||
def build_browser_args(
|
||||
proxy_location: ProxyLocation | None = None,
|
||||
cdp_port: int | None = None,
|
||||
extra_http_headers: dict[str, str] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
video_dir = f"{settings.VIDEO_PATH}/{datetime.utcnow().strftime('%Y-%m-%d')}"
|
||||
har_dir = (
|
||||
f"{settings.HAR_PATH}/{datetime.utcnow().strftime('%Y-%m-%d')}/{BrowserContextFactory.get_subdir()}.har"
|
||||
@@ -214,6 +218,7 @@ class BrowserContextFactory:
|
||||
"width": settings.BROWSER_WIDTH,
|
||||
"height": settings.BROWSER_HEIGHT,
|
||||
},
|
||||
"extra_http_headers": extra_http_headers,
|
||||
}
|
||||
|
||||
if settings.ENABLE_PROXY:
|
||||
@@ -405,7 +410,10 @@ def _is_chrome_running() -> bool:
|
||||
|
||||
|
||||
async def _create_headless_chromium(
|
||||
playwright: Playwright, proxy_location: ProxyLocation | None = None, **kwargs: dict
|
||||
playwright: Playwright,
|
||||
proxy_location: ProxyLocation | None = None,
|
||||
extra_http_headers: dict[str, str] | None = None,
|
||||
**kwargs: dict,
|
||||
) -> tuple[BrowserContext, BrowserArtifacts, BrowserCleanupFunc]:
|
||||
user_data_dir = make_temp_directory(prefix="skyvern_browser_")
|
||||
download_dir = initialize_download_dir()
|
||||
@@ -414,7 +422,9 @@ async def _create_headless_chromium(
|
||||
download_dir=download_dir,
|
||||
)
|
||||
cdp_port: int | None = _get_cdp_port(kwargs)
|
||||
browser_args = BrowserContextFactory.build_browser_args(proxy_location=proxy_location, cdp_port=cdp_port)
|
||||
browser_args = BrowserContextFactory.build_browser_args(
|
||||
proxy_location=proxy_location, cdp_port=cdp_port, extra_http_headers=extra_http_headers
|
||||
)
|
||||
browser_args.update(
|
||||
{
|
||||
"user_data_dir": user_data_dir,
|
||||
@@ -428,7 +438,10 @@ async def _create_headless_chromium(
|
||||
|
||||
|
||||
async def _create_headful_chromium(
|
||||
playwright: Playwright, proxy_location: ProxyLocation | None = None, **kwargs: dict
|
||||
playwright: Playwright,
|
||||
proxy_location: ProxyLocation | None = None,
|
||||
extra_http_headers: dict[str, str] | None = None,
|
||||
**kwargs: dict,
|
||||
) -> tuple[BrowserContext, BrowserArtifacts, BrowserCleanupFunc]:
|
||||
user_data_dir = make_temp_directory(prefix="skyvern_browser_")
|
||||
download_dir = initialize_download_dir()
|
||||
@@ -437,7 +450,9 @@ async def _create_headful_chromium(
|
||||
download_dir=download_dir,
|
||||
)
|
||||
cdp_port: int | None = _get_cdp_port(kwargs)
|
||||
browser_args = BrowserContextFactory.build_browser_args(proxy_location=proxy_location, cdp_port=cdp_port)
|
||||
browser_args = BrowserContextFactory.build_browser_args(
|
||||
proxy_location=proxy_location, cdp_port=cdp_port, extra_http_headers=extra_http_headers
|
||||
)
|
||||
browser_args.update(
|
||||
{
|
||||
"user_data_dir": user_data_dir,
|
||||
@@ -479,7 +494,10 @@ def is_valid_chromium_user_data_dir(directory: str) -> bool:
|
||||
|
||||
|
||||
async def _create_cdp_connection_browser(
|
||||
playwright: Playwright, proxy_location: ProxyLocation | None = None, **kwargs: dict
|
||||
playwright: Playwright,
|
||||
proxy_location: ProxyLocation | None = None,
|
||||
extra_http_headers: dict[str, str] | None = None,
|
||||
**kwargs: dict,
|
||||
) -> tuple[BrowserContext, BrowserArtifacts, BrowserCleanupFunc]:
|
||||
browser_type = settings.BROWSER_TYPE
|
||||
browser_path = settings.CHROME_EXECUTABLE_PATH
|
||||
@@ -528,7 +546,7 @@ async def _create_cdp_connection_browser(
|
||||
else:
|
||||
LOG.info("Port 9222 is in use, using existing browser")
|
||||
|
||||
browser_args = BrowserContextFactory.build_browser_args()
|
||||
browser_args = BrowserContextFactory.build_browser_args(extra_http_headers=extra_http_headers)
|
||||
|
||||
browser_artifacts = BrowserContextFactory.build_browser_artifacts(
|
||||
har_path=browser_args["record_har_path"],
|
||||
@@ -549,6 +567,7 @@ async def _create_cdp_connection_browser(
|
||||
browser_context = await browser.new_context(
|
||||
record_video_dir=browser_args["record_video_dir"],
|
||||
viewport=browser_args["viewport"],
|
||||
extra_http_headers=browser_args["extra_http_headers"],
|
||||
)
|
||||
LOG.info(
|
||||
"Launched browser CDP connection",
|
||||
@@ -608,6 +627,7 @@ class BrowserState:
|
||||
task_id: str | None = None,
|
||||
workflow_run_id: str | None = None,
|
||||
organization_id: str | None = None,
|
||||
extra_http_headers: dict[str, str] | None = None,
|
||||
) -> None:
|
||||
if self.browser_context is None:
|
||||
LOG.info("creating browser context")
|
||||
@@ -622,6 +642,7 @@ class BrowserState:
|
||||
task_id=task_id,
|
||||
workflow_run_id=workflow_run_id,
|
||||
organization_id=organization_id,
|
||||
extra_http_headers=extra_http_headers,
|
||||
)
|
||||
self.browser_context = browser_context
|
||||
self.browser_artifacts = browser_artifacts
|
||||
@@ -752,6 +773,7 @@ class BrowserState:
|
||||
task_id: str | None = None,
|
||||
workflow_run_id: str | None = None,
|
||||
organization_id: str | None = None,
|
||||
extra_http_headers: dict[str, str] | None = None,
|
||||
) -> Page:
|
||||
page = await self.get_working_page()
|
||||
if page is not None:
|
||||
@@ -764,6 +786,7 @@ class BrowserState:
|
||||
task_id=task_id,
|
||||
workflow_run_id=workflow_run_id,
|
||||
organization_id=organization_id,
|
||||
extra_http_headers=extra_http_headers,
|
||||
)
|
||||
except Exception as e:
|
||||
error_message = str(e)
|
||||
@@ -778,6 +801,7 @@ class BrowserState:
|
||||
task_id=task_id,
|
||||
workflow_run_id=workflow_run_id,
|
||||
organization_id=organization_id,
|
||||
extra_http_headers=extra_http_headers,
|
||||
)
|
||||
page = await self.__assert_page()
|
||||
|
||||
@@ -791,6 +815,7 @@ class BrowserState:
|
||||
task_id=task_id,
|
||||
workflow_run_id=workflow_run_id,
|
||||
organization_id=organization_id,
|
||||
extra_http_headers=extra_http_headers,
|
||||
)
|
||||
page = await self.__assert_page()
|
||||
return page
|
||||
|
||||
@@ -31,6 +31,7 @@ class BrowserManager:
|
||||
task_id: str | None = None,
|
||||
workflow_run_id: str | None = None,
|
||||
organization_id: str | None = None,
|
||||
extra_http_headers: dict[str, str] | None = None,
|
||||
) -> BrowserState:
|
||||
pw = await async_playwright().start()
|
||||
(
|
||||
@@ -44,6 +45,7 @@ class BrowserManager:
|
||||
task_id=task_id,
|
||||
workflow_run_id=workflow_run_id,
|
||||
organization_id=organization_id,
|
||||
extra_http_headers=extra_http_headers,
|
||||
)
|
||||
return BrowserState(
|
||||
pw=pw,
|
||||
@@ -109,6 +111,7 @@ class BrowserManager:
|
||||
url=task.url,
|
||||
task_id=task.task_id,
|
||||
organization_id=task.organization_id,
|
||||
extra_http_headers=task.extra_http_headers,
|
||||
)
|
||||
|
||||
self.pages[task.task_id] = browser_state
|
||||
@@ -118,7 +121,11 @@ class BrowserManager:
|
||||
# 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, organization_id=task.organization_id
|
||||
url=task.url,
|
||||
proxy_location=task.proxy_location,
|
||||
task_id=task.task_id,
|
||||
organization_id=task.organization_id,
|
||||
extra_http_headers=task.extra_http_headers,
|
||||
)
|
||||
return browser_state
|
||||
|
||||
@@ -173,6 +180,7 @@ class BrowserManager:
|
||||
url=url,
|
||||
workflow_run_id=workflow_run.workflow_run_id,
|
||||
organization_id=workflow_run.organization_id,
|
||||
extra_http_headers=workflow_run.extra_http_headers,
|
||||
)
|
||||
|
||||
self.pages[workflow_run_id] = browser_state
|
||||
@@ -186,6 +194,7 @@ class BrowserManager:
|
||||
proxy_location=workflow_run.proxy_location,
|
||||
workflow_run_id=workflow_run.workflow_run_id,
|
||||
organization_id=workflow_run.organization_id,
|
||||
extra_http_headers=workflow_run.extra_http_headers,
|
||||
)
|
||||
return browser_state
|
||||
|
||||
|
||||
Reference in New Issue
Block a user