From 965150b45b25d5667b0eb9f2a37d74ce93450833 Mon Sep 17 00:00:00 2001 From: Andrew Neilson Date: Fri, 22 May 2026 14:31:42 -0400 Subject: [PATCH] Re-introduce cdp_connect_headers (#11067) with NOWAIT migration (#11198) (#6117) --- ...55-9f512f2da31e_add_cdp_connect_headers.py | 100 +++++++++++ skyvern/forge/agent.py | 2 + skyvern/forge/sdk/db/models.py | 4 + skyvern/forge/sdk/db/repositories/observer.py | 2 + skyvern/forge/sdk/db/repositories/tasks.py | 2 + .../sdk/db/repositories/workflow_runs.py | 5 + .../forge/sdk/db/repositories/workflows.py | 8 + skyvern/forge/sdk/db/utils.py | 3 + skyvern/forge/sdk/routes/agent_protocol.py | 5 + skyvern/forge/sdk/schemas/task_v2.py | 9 +- skyvern/forge/sdk/schemas/tasks.py | 15 +- skyvern/forge/sdk/workflow/models/workflow.py | 19 +- skyvern/forge/sdk/workflow/service.py | 36 ++++ skyvern/schemas/run_blocks.py | 7 + skyvern/schemas/runs.py | 29 ++- skyvern/schemas/workflows.py | 8 +- skyvern/services/task_v2_service.py | 4 + skyvern/utils/secret_headers.py | 39 ++++ skyvern/webeye/browser_factory.py | 26 ++- skyvern/webeye/browser_state.py | 2 + skyvern/webeye/real_browser_manager.py | 6 + skyvern/webeye/real_browser_state.py | 6 + .../test_cdp_connect_headers_migration.py | 170 ++++++++++++++++++ tests/unit/test_cdp_connection.py | 32 ++++ .../unit/test_workflow_browser_profile_id.py | 15 ++ .../test_workflow_run_parameter_batching.py | 1 + 26 files changed, 548 insertions(+), 7 deletions(-) create mode 100644 alembic/versions/2026_05_22_1755-9f512f2da31e_add_cdp_connect_headers.py create mode 100644 skyvern/utils/secret_headers.py create mode 100644 tests/unit/test_cdp_connect_headers_migration.py diff --git a/alembic/versions/2026_05_22_1755-9f512f2da31e_add_cdp_connect_headers.py b/alembic/versions/2026_05_22_1755-9f512f2da31e_add_cdp_connect_headers.py new file mode 100644 index 000000000..58d87d73d --- /dev/null +++ b/alembic/versions/2026_05_22_1755-9f512f2da31e_add_cdp_connect_headers.py @@ -0,0 +1,100 @@ +"""add cdp_connect_headers + +Revision ID: 9f512f2da31e +Revises: 11a965fb5d82 +Create Date: 2026-05-22T17:55:29.232754+00:00 + +""" + +import time +from typing import Sequence, Union + +import sqlalchemy as sa +from sqlalchemy.exc import DBAPIError + +from alembic import op + +# revision identifiers, used by Alembic. +revision: str = "9f512f2da31e" +down_revision: Union[str, None] = "11a965fb5d82" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + +_TABLES = ("observer_cruises", "tasks", "workflow_runs", "workflows") +_COLUMN = "cdp_connect_headers" +_STATEMENT_TIMEOUT = "5s" +_LOCK_NOT_AVAILABLE_SQLSTATE = "55P03" +_MIGRATION_RETRY_SECONDS = 10 * 60 +_BACKOFF_SECONDS = 5 + + +def _is_lock_not_available(exc: DBAPIError) -> bool: + orig = getattr(exc, "orig", None) + return ( + getattr(orig, "sqlstate", None) == _LOCK_NOT_AVAILABLE_SQLSTATE + or getattr(orig, "pgcode", None) == _LOCK_NOT_AVAILABLE_SQLSTATE + ) + + +def _column_exists(table: str) -> bool: + result = op.get_bind().execute( + sa.text( + """ + SELECT 1 + FROM information_schema.columns + WHERE table_schema = current_schema() + AND table_name = :table_name + AND column_name = :column_name + """ + ), + {"table_name": table, "column_name": _COLUMN}, + ) + return result.first() is not None + + +def _execute_transactional_schema_change(table: str, statement: str) -> None: + try: + op.execute("BEGIN") + op.execute(f"SET LOCAL statement_timeout = '{_STATEMENT_TIMEOUT}'") + op.execute(f'LOCK TABLE "{table}" IN ACCESS EXCLUSIVE MODE NOWAIT') + op.execute(statement) + op.execute("COMMIT") + except Exception as exc: + try: + op.execute("ROLLBACK") + except Exception as rollback_exc: + raise RuntimeError( + f"rollback failed after schema change error; aborting migration: {rollback_exc}" + ) from exc + raise + + +def _execute_with_retry(table: str, statement: str, deadline: float) -> None: + while True: + try: + _execute_transactional_schema_change(table, statement) + return + except DBAPIError as exc: + if not _is_lock_not_available(exc) or time.monotonic() >= deadline: + raise + time.sleep(_BACKOFF_SECONDS) + + +def upgrade() -> None: + deadline = time.monotonic() + _MIGRATION_RETRY_SECONDS + with op.get_context().autocommit_block(): + for table in reversed(_TABLES): + if not _column_exists(table): + _execute_with_retry( + table, + f'ALTER TABLE "{table}" ADD COLUMN IF NOT EXISTS {_COLUMN} JSON', + deadline, + ) + + +def downgrade() -> None: + deadline = time.monotonic() + _MIGRATION_RETRY_SECONDS + with op.get_context().autocommit_block(): + for table in _TABLES: + if _column_exists(table): + _execute_with_retry(table, f'ALTER TABLE "{table}" DROP COLUMN IF EXISTS {_COLUMN}', deadline) diff --git a/skyvern/forge/agent.py b/skyvern/forge/agent.py index 2724ce72f..1de1c185a 100644 --- a/skyvern/forge/agent.py +++ b/skyvern/forge/agent.py @@ -500,6 +500,7 @@ class ForgeAgent: model=task_block.model, max_screenshot_scrolling_times=workflow_run.max_screenshot_scrolls, extra_http_headers=workflow_run.extra_http_headers, + cdp_connect_headers=workflow_run.cdp_connect_headers, browser_address=workflow_run.browser_address, browser_session_id=workflow_run.browser_session_id, download_timeout=task_block.download_timeout, @@ -571,6 +572,7 @@ class ForgeAgent: model=task_request.model, max_screenshot_scrolling_times=task_request.max_screenshot_scrolls, extra_http_headers=task_request.extra_http_headers, + cdp_connect_headers=task_request.cdp_connect_headers, browser_session_id=task_request.browser_session_id, browser_address=task_request.browser_address, include_extracted_text=task_request.include_extracted_text, diff --git a/skyvern/forge/sdk/db/models.py b/skyvern/forge/sdk/db/models.py index d6fc532f5..a6e4be6b0 100644 --- a/skyvern/forge/sdk/db/models.py +++ b/skyvern/forge/sdk/db/models.py @@ -109,6 +109,7 @@ class TaskModel(Base): proxy_location = Column(String) extracted_information_schema = Column(JSON) extra_http_headers = Column(JSON, nullable=True) + cdp_connect_headers = Column(JSON, nullable=True) workflow_run_id = Column(String, ForeignKey("workflow_runs.workflow_run_id"), index=True) order = Column(Integer, nullable=True) retry = Column(Integer, nullable=True) @@ -326,6 +327,7 @@ class WorkflowModel(SoftDeleteMixin, Base): webhook_callback_url = Column(String) max_screenshot_scrolling_times = Column(Integer, nullable=True) extra_http_headers = Column(JSON, nullable=True) + cdp_connect_headers = Column(JSON, nullable=True) totp_verification_url = Column(String) totp_identifier = Column(String) persist_browser_session = Column(Boolean, default=False, nullable=False) @@ -444,6 +446,7 @@ class WorkflowRunModel(Base): totp_identifier = Column(String) max_screenshot_scrolling_times = Column(Integer, nullable=True) extra_http_headers = Column(JSON, nullable=True) + cdp_connect_headers = Column(JSON, nullable=True) browser_address = Column(String, nullable=True, index=True) script_run = Column(JSON, nullable=True) job_id = Column(String, nullable=True, index=True) @@ -913,6 +916,7 @@ class TaskV2Model(Base): max_steps = Column(Integer, nullable=True) max_screenshot_scrolling_times = Column(Integer, nullable=True) extra_http_headers = Column(JSON, nullable=True) + cdp_connect_headers = Column(JSON, nullable=True) browser_address = Column(String, nullable=True) generate_script = Column(Boolean, default=False, nullable=False) run_with = Column(String, nullable=True) # 'agent' or 'code' diff --git a/skyvern/forge/sdk/db/repositories/observer.py b/skyvern/forge/sdk/db/repositories/observer.py index af3a34a93..e0f2bd633 100644 --- a/skyvern/forge/sdk/db/repositories/observer.py +++ b/skyvern/forge/sdk/db/repositories/observer.py @@ -202,6 +202,7 @@ class ObserverRepository(BaseRepository): model: dict[str, Any] | None = None, max_screenshot_scrolling_times: int | None = None, extra_http_headers: dict[str, str] | None = None, + cdp_connect_headers: dict[str, str] | None = None, browser_address: str | None = None, run_with: str | None = None, ) -> TaskV2: @@ -225,6 +226,7 @@ class ObserverRepository(BaseRepository): model=model, max_screenshot_scrolling_times=max_screenshot_scrolling_times, extra_http_headers=extra_http_headers, + cdp_connect_headers=cdp_connect_headers, browser_address=browser_address, run_with=run_with, ) diff --git a/skyvern/forge/sdk/db/repositories/tasks.py b/skyvern/forge/sdk/db/repositories/tasks.py index 6424b75e4..c2341f219 100644 --- a/skyvern/forge/sdk/db/repositories/tasks.py +++ b/skyvern/forge/sdk/db/repositories/tasks.py @@ -64,6 +64,7 @@ class TasksRepository(BaseRepository): model: dict[str, Any] | None = None, max_screenshot_scrolling_times: int | None = None, extra_http_headers: dict[str, str] | None = None, + cdp_connect_headers: dict[str, str] | None = None, browser_session_id: str | None = None, browser_address: str | None = None, download_timeout: float | None = None, @@ -110,6 +111,7 @@ class TasksRepository(BaseRepository): model=model, max_screenshot_scrolling_times=max_screenshot_scrolling_times, extra_http_headers=extra_http_headers, + cdp_connect_headers=cdp_connect_headers, browser_session_id=browser_session_id, browser_address=browser_address, download_timeout=download_timeout, diff --git a/skyvern/forge/sdk/db/repositories/workflow_runs.py b/skyvern/forge/sdk/db/repositories/workflow_runs.py index 1bad412e0..05d42f95f 100644 --- a/skyvern/forge/sdk/db/repositories/workflow_runs.py +++ b/skyvern/forge/sdk/db/repositories/workflow_runs.py @@ -155,6 +155,7 @@ class WorkflowRunsRepository(BaseRepository): parent_workflow_run_id: str | None = None, max_screenshot_scrolling_times: int | None = None, extra_http_headers: dict[str, str] | None = None, + cdp_connect_headers: dict[str, str] | None = None, browser_address: str | None = None, sequential_key: str | None = None, run_with: str | None = None, @@ -185,6 +186,7 @@ class WorkflowRunsRepository(BaseRepository): parent_workflow_run_id=parent_workflow_run_id, max_screenshot_scrolling_times=max_screenshot_scrolling_times, extra_http_headers=extra_http_headers, + cdp_connect_headers=cdp_connect_headers, browser_address=browser_address, sequential_key=sequential_key, run_with=run_with, @@ -224,6 +226,7 @@ class WorkflowRunsRepository(BaseRepository): browser_profile_id: str | None | object = _UNSET, browser_address: str | None = None, extra_http_headers: dict[str, str] | None = None, + cdp_connect_headers: dict[str, str] | None = None, failure_category: list[dict[str, Any]] | None = None, started_at: datetime | None | object = _UNSET, queued_at: datetime | None | object = _UNSET, @@ -269,6 +272,8 @@ class WorkflowRunsRepository(BaseRepository): workflow_run.browser_address = browser_address if extra_http_headers is not None: workflow_run.extra_http_headers = extra_http_headers + if cdp_connect_headers is not None: + workflow_run.cdp_connect_headers = cdp_connect_headers # 2FA verification code waiting state updates if waiting_for_verification_code is not None: workflow_run.waiting_for_verification_code = waiting_for_verification_code diff --git a/skyvern/forge/sdk/db/repositories/workflows.py b/skyvern/forge/sdk/db/repositories/workflows.py index 7115d45d8..51791cb7d 100644 --- a/skyvern/forge/sdk/db/repositories/workflows.py +++ b/skyvern/forge/sdk/db/repositories/workflows.py @@ -96,6 +96,7 @@ class WorkflowsRepository(BaseRepository): webhook_callback_url: str | None = None, max_screenshot_scrolling_times: int | None = None, extra_http_headers: dict[str, str] | None = None, + cdp_connect_headers: dict[str, str] | None = None, totp_verification_url: str | None = None, totp_identifier: str | None = None, persist_browser_session: bool = False, @@ -129,6 +130,7 @@ class WorkflowsRepository(BaseRepository): totp_identifier=totp_identifier, max_screenshot_scrolling_times=max_screenshot_scrolling_times, extra_http_headers=extra_http_headers, + cdp_connect_headers=cdp_connect_headers, persist_browser_session=persist_browser_session, browser_profile_id=browser_profile_id, model=model, @@ -612,6 +614,7 @@ class WorkflowsRepository(BaseRepository): model: dict[str, Any] | None | object = _UNSET, max_screenshot_scrolling_times: int | None | object = _UNSET, extra_http_headers: dict[str, str] | None | object = _UNSET, + cdp_connect_headers: dict[str, str] | None | object = _UNSET, ai_fallback: bool | None = None, run_sequentially: bool | None = None, sequential_key: str | None | object = _UNSET, @@ -657,6 +660,8 @@ class WorkflowsRepository(BaseRepository): workflow.max_screenshot_scrolling_times = max_screenshot_scrolling_times if extra_http_headers is not _UNSET: workflow.extra_http_headers = extra_http_headers + if cdp_connect_headers is not _UNSET: + workflow.cdp_connect_headers = cdp_connect_headers if ai_fallback is not None: workflow.ai_fallback = ai_fallback if run_sequentially is not None: @@ -906,6 +911,7 @@ class WorkflowsRepository(BaseRepository): model: dict[str, Any] | None | object = _UNSET, max_screenshot_scrolling_times: int | None | object = _UNSET, extra_http_headers: dict[str, str] | None | object = _UNSET, + cdp_connect_headers: dict[str, str] | None | object = _UNSET, ai_fallback: bool | None = None, run_sequentially: bool | None = None, sequential_key: str | None | object = _UNSET, @@ -985,6 +991,8 @@ class WorkflowsRepository(BaseRepository): workflow.max_screenshot_scrolling_times = max_screenshot_scrolling_times if extra_http_headers is not _UNSET: workflow.extra_http_headers = extra_http_headers + if cdp_connect_headers is not _UNSET: + workflow.cdp_connect_headers = cdp_connect_headers if ai_fallback is not None: workflow.ai_fallback = ai_fallback if run_sequentially is not None: diff --git a/skyvern/forge/sdk/db/utils.py b/skyvern/forge/sdk/db/utils.py index f2980beb8..26db9b6ab 100644 --- a/skyvern/forge/sdk/db/utils.py +++ b/skyvern/forge/sdk/db/utils.py @@ -310,6 +310,7 @@ def convert_to_task(task_obj: TaskModel, debug_enabled: bool = False, workflow_p proxy_location=deserialize_proxy_location(task_obj.proxy_location), extracted_information_schema=task_obj.extracted_information_schema, extra_http_headers=task_obj.extra_http_headers, + cdp_connect_headers=task_obj.cdp_connect_headers, workflow_run_id=task_obj.workflow_run_id, workflow_permanent_id=workflow_permanent_id, order=task_obj.order, @@ -507,6 +508,7 @@ def convert_to_workflow( deleted_at=workflow_model.deleted_at, status=WorkflowStatus(workflow_model.status), extra_http_headers=workflow_model.extra_http_headers, + cdp_connect_headers=workflow_model.cdp_connect_headers, run_with=workflow_model.run_with, ai_fallback=workflow_model.ai_fallback, cache_key=workflow_model.cache_key, @@ -555,6 +557,7 @@ def convert_to_workflow_run( workflow_title=workflow_title, max_screenshot_scrolls=workflow_run_model.max_screenshot_scrolling_times, extra_http_headers=workflow_run_model.extra_http_headers, + cdp_connect_headers=workflow_run_model.cdp_connect_headers, browser_address=workflow_run_model.browser_address, job_id=workflow_run_model.job_id, depends_on_workflow_run_id=workflow_run_model.depends_on_workflow_run_id, diff --git a/skyvern/forge/sdk/routes/agent_protocol.py b/skyvern/forge/sdk/routes/agent_protocol.py index fd9b2c8bc..7912ecec6 100644 --- a/skyvern/forge/sdk/routes/agent_protocol.py +++ b/skyvern/forge/sdk/routes/agent_protocol.py @@ -263,6 +263,7 @@ async def run_task( model=run_request.model, max_screenshot_scrolls=run_request.max_screenshot_scrolls, extra_http_headers=run_request.extra_http_headers, + cdp_connect_headers=run_request.cdp_connect_headers, browser_address=run_request.browser_address, ) task_v1_response = await task_v1_service.run_task( @@ -330,6 +331,7 @@ async def run_task( model=run_request.model, max_screenshot_scrolling_times=run_request.max_screenshot_scrolls, extra_http_headers=run_request.extra_http_headers, + cdp_connect_headers=run_request.cdp_connect_headers, browser_session_id=run_request.browser_session_id, browser_address=run_request.browser_address, run_with=run_request.run_with, @@ -439,6 +441,7 @@ async def run_workflow( browser_profile_id=workflow_run_request.browser_profile_id, max_screenshot_scrolls=workflow_run_request.max_screenshot_scrolls, extra_http_headers=workflow_run_request.extra_http_headers, + cdp_connect_headers=workflow_run_request.cdp_connect_headers, browser_address=workflow_run_request.browser_address, run_with=workflow_run_request.run_with, ai_fallback=workflow_run_request.ai_fallback, @@ -745,6 +748,7 @@ async def create_workflow_from_prompt( proxy_location=request.proxy_location, max_screenshot_scrolling_times=request.max_screenshot_scrolls, extra_http_headers=request.extra_http_headers, + cdp_connect_headers=request.cdp_connect_headers, max_iterations=x_max_iterations_override, max_steps=x_max_steps_override, status=WorkflowStatus.published if request.publish_workflow else WorkflowStatus.auto_generated, @@ -3553,6 +3557,7 @@ async def run_task_v2( max_screenshot_scrolling_times=data.max_screenshot_scrolls, browser_session_id=data.browser_session_id, extra_http_headers=data.extra_http_headers, + cdp_connect_headers=data.cdp_connect_headers, browser_address=data.browser_address, trigger_type=legacy_v2_trigger_type, ) diff --git a/skyvern/forge/sdk/schemas/task_v2.py b/skyvern/forge/sdk/schemas/task_v2.py index c839c752b..6c5ab1101 100644 --- a/skyvern/forge/sdk/schemas/task_v2.py +++ b/skyvern/forge/sdk/schemas/task_v2.py @@ -3,10 +3,11 @@ from datetime import datetime from enum import StrEnum from typing import Any -from pydantic import BaseModel, ConfigDict, Field, field_validator +from pydantic import BaseModel, ConfigDict, Field, field_serializer, field_validator from skyvern.forge.sdk.settings_manager import SettingsManager from skyvern.schemas.runs import GeoTarget, ProxyLocation, ProxyLocationInput +from skyvern.utils.secret_headers import mask_header_values from skyvern.utils.url_validators import validate_url DEFAULT_WORKFLOW_TITLE = "New Workflow" @@ -53,6 +54,7 @@ class TaskV2(BaseModel): finished_at: datetime | None = None max_screenshot_scrolls: int | None = Field(default=None, alias="max_screenshot_scrolling_times") extra_http_headers: dict[str, str] | None = None + cdp_connect_headers: dict[str, str] | None = None browser_address: str | None = None run_with: str | None = None failure_category: list[dict[str, Any]] | None = None @@ -117,6 +119,10 @@ class TaskV2(BaseModel): def deserialize_proxy_location(cls, proxy_location: ProxyLocationInput | str) -> ProxyLocationInput: return cls._parse_proxy_location(proxy_location) + @field_serializer("cdp_connect_headers") + def _mask_cdp_connect_headers(self, headers: dict[str, str] | None) -> dict[str, str] | None: + return mask_header_values(headers) + class ThoughtType(StrEnum): """ @@ -207,6 +213,7 @@ class TaskV2Request(BaseModel): workflow_system_prompt: str | None = None max_screenshot_scrolls: int | None = None extra_http_headers: dict[str, str] | None = None + cdp_connect_headers: dict[str, str] | None = None browser_address: str | None = None run_with: str | None = None ai_fallback: bool = False diff --git a/skyvern/forge/sdk/schemas/tasks.py b/skyvern/forge/sdk/schemas/tasks.py index d00061999..fd64942a3 100644 --- a/skyvern/forge/sdk/schemas/tasks.py +++ b/skyvern/forge/sdk/schemas/tasks.py @@ -5,7 +5,7 @@ from enum import StrEnum from typing import Any from fastapi import status -from pydantic import BaseModel, Field, field_validator, model_validator +from pydantic import BaseModel, Field, field_serializer, field_validator, model_validator from typing_extensions import Self from skyvern.exceptions import ( @@ -20,6 +20,7 @@ from skyvern.forge.sdk.settings_manager import SettingsManager from skyvern.schemas.docs.doc_strings import PROXY_LOCATION_DOC_STRING from skyvern.schemas.runs import ProxyLocationInput from skyvern.utils.prompt_truncation import EXTRACTION_GOAL_MAX_TOKENS +from skyvern.utils.secret_headers import mask_header_values from skyvern.utils.token_counter import count_tokens from skyvern.utils.url_validators import validate_url @@ -87,6 +88,14 @@ class TaskBase(BaseModel): extra_http_headers: dict[str, str] | None = Field( None, description="The extra HTTP headers for the requests in browser." ) + cdp_connect_headers: dict[str, str] | None = Field( + None, + description=( + "HTTP headers attached ONLY to the CDP WebSocket handshake when connecting to " + "a remote browser via browser_address. Use this for browser-provider auth. " + "Never forwarded to target websites." + ), + ) complete_criterion: str | None = Field( default=None, description="Criterion to complete", examples=["Complete if 'hello world' shows up on the page"] ) @@ -130,6 +139,10 @@ class TaskBase(BaseModel): description="If False, omit the scraped page text dump from the extract-information prompt. ExtractionBlock opts out; everything else keeps the default.", ) + @field_serializer("cdp_connect_headers") + def _mask_cdp_connect_headers(self, headers: dict[str, str] | None) -> dict[str, str] | None: + return mask_header_values(headers) + class TaskRequest(TaskBase): url: str = Field( diff --git a/skyvern/forge/sdk/workflow/models/workflow.py b/skyvern/forge/sdk/workflow/models/workflow.py index dad401140..b5f371270 100644 --- a/skyvern/forge/sdk/workflow/models/workflow.py +++ b/skyvern/forge/sdk/workflow/models/workflow.py @@ -2,7 +2,7 @@ from datetime import datetime from enum import StrEnum from typing import Any, List -from pydantic import BaseModel, Field, field_validator +from pydantic import BaseModel, Field, field_serializer, field_validator from typing_extensions import deprecated from skyvern.forge.sdk.db.enums import WorkflowRunTriggerType @@ -18,6 +18,7 @@ from skyvern.forge.sdk.workflow.models.parameter import PARAMETER_TYPE, OutputPa from skyvern.forge.sdk.workflow.models.validators import normalize_run_metadata, normalize_run_with from skyvern.schemas.runs import ProxyLocationInput, ScriptRunResponse from skyvern.schemas.workflows import BlockType, WorkflowStatus +from skyvern.utils.secret_headers import mask_header_values from skyvern.utils.url_validators import validate_url @@ -32,6 +33,7 @@ class WorkflowRequestBody(BaseModel): browser_profile_id: str | None = None max_screenshot_scrolls: int | None = None extra_http_headers: dict[str, str] | None = None + cdp_connect_headers: dict[str, str] | None = None browser_address: str | None = None run_with: str | None = None ai_fallback: bool | None = None @@ -115,6 +117,7 @@ class Workflow(BaseModel): status: WorkflowStatus = WorkflowStatus.published max_screenshot_scrolls: int | None = None extra_http_headers: dict[str, str] | None = None + cdp_connect_headers: dict[str, str] | None = None run_with: str = "agent" ai_fallback: bool = True cache_key: str | None = None @@ -133,6 +136,10 @@ class Workflow(BaseModel): def _normalize_run_with(cls, v: str | None) -> str: return normalize_run_with(v) + @field_serializer("cdp_connect_headers") + def _mask_cdp_connect_headers(self, headers: dict[str, str] | None) -> dict[str, str] | None: + return mask_header_values(headers) + created_at: datetime modified_at: datetime deleted_at: datetime | None = None @@ -195,6 +202,7 @@ class WorkflowRun(BaseModel): debug_session_id: str | None = None status: WorkflowRunStatus extra_http_headers: dict[str, str] | None = None + cdp_connect_headers: dict[str, str] | None = None proxy_location: ProxyLocationInput = None webhook_callback_url: str | None = None webhook_failure_reason: str | None = None @@ -228,6 +236,10 @@ class WorkflowRun(BaseModel): return None return normalize_run_with(v) + @field_serializer("cdp_connect_headers") + def _mask_cdp_connect_headers(self, headers: dict[str, str] | None) -> dict[str, str] | None: + return mask_header_values(headers) + queued_at: datetime | None = None started_at: datetime | None = None finished_at: datetime | None = None @@ -299,6 +311,7 @@ class WorkflowRunResponseBase(BaseModel): totp_verification_url: str | None = None totp_identifier: str | None = None extra_http_headers: dict[str, str] | None = None + cdp_connect_headers: dict[str, str] | None = None queued_at: datetime | None = None started_at: datetime | None = None finished_at: datetime | None = None @@ -335,6 +348,10 @@ class WorkflowRunResponseBase(BaseModel): def _normalize_run_with(cls, v: str | None) -> str: return normalize_run_with(v) + @field_serializer("cdp_connect_headers") + def _mask_cdp_connect_headers(self, headers: dict[str, str] | None) -> dict[str, str] | None: + return mask_header_values(headers) + class WorkflowRunWithWorkflowResponse(WorkflowRunResponseBase): workflow: Workflow diff --git a/skyvern/forge/sdk/workflow/service.py b/skyvern/forge/sdk/workflow/service.py index b36c06672..effbd3e66 100644 --- a/skyvern/forge/sdk/workflow/service.py +++ b/skyvern/forge/sdk/workflow/service.py @@ -130,6 +130,7 @@ from skyvern.schemas.workflows import ( from skyvern.services import script_service, workflow_script_service from skyvern.services.webhook_delivery import PreparedWorkflowWebhook, deliver_webhook_with_retries from skyvern.utils.css_selector import build_action_summaries_with_timing # shared with script_service +from skyvern.utils.secret_headers import merge_masked_headers from skyvern.utils.url_validators import validate_url as validate_url_with_blocked_host_check from skyvern.webeye.browser_state import BrowserState @@ -902,6 +903,13 @@ class WorkflowService: and workflow.browser_profile_id is not None ): workflow_request.browser_profile_id = workflow.browser_profile_id + if workflow_request.cdp_connect_headers is None: + if workflow.cdp_connect_headers is not None: + workflow_request.cdp_connect_headers = workflow.cdp_connect_headers + else: + workflow_request.cdp_connect_headers = merge_masked_headers( + workflow_request.cdp_connect_headers, workflow.cdp_connect_headers + ) if workflow_request.run_with is None: workflow_request.run_with = workflow.run_with @@ -3506,6 +3514,7 @@ class WorkflowService: is_saved_task: bool = False, status: WorkflowStatus = WorkflowStatus.published, extra_http_headers: dict[str, str] | None = None, + cdp_connect_headers: dict[str, str] | None = None, run_with: str | None = None, cache_key: str | None = None, ai_fallback: bool | None = None, @@ -3537,6 +3546,7 @@ class WorkflowService: is_saved_task=is_saved_task, status=status, extra_http_headers=extra_http_headers, + cdp_connect_headers=cdp_connect_headers, run_with=run_with, cache_key=cache_key, ai_fallback=True if ai_fallback is None else ai_fallback, @@ -3564,6 +3574,7 @@ class WorkflowService: proxy_location: ProxyLocationInput = None, max_screenshot_scrolling_times: int | None = None, extra_http_headers: dict[str, str] | None = None, + cdp_connect_headers: dict[str, str] | None = None, max_iterations: int | None = None, max_steps: int | None = None, status: WorkflowStatus = WorkflowStatus.auto_generated, @@ -3704,6 +3715,7 @@ class WorkflowService: totp_identifier=totp_identifier, max_screenshot_scrolling_times=max_screenshot_scrolling_times, extra_http_headers=extra_http_headers, + cdp_connect_headers=cdp_connect_headers, status=status, run_with=run_with, ai_fallback=ai_fallback, @@ -3916,6 +3928,7 @@ class WorkflowService: model: dict[str, Any] | None | object = _UNSET, max_screenshot_scrolling_times: int | None | object = _UNSET, extra_http_headers: dict[str, str] | None | object = _UNSET, + cdp_connect_headers: dict[str, str] | None | object = _UNSET, run_with: str | None = None, ai_fallback: bool | None = None, cache_key: str | None = None, @@ -3938,6 +3951,7 @@ class WorkflowService: model=model, max_screenshot_scrolling_times=max_screenshot_scrolling_times, extra_http_headers=extra_http_headers, + cdp_connect_headers=cdp_connect_headers, run_with=run_with, ai_fallback=ai_fallback, cache_key=cache_key, @@ -3961,6 +3975,7 @@ class WorkflowService: model=model, max_screenshot_scrolling_times=max_screenshot_scrolling_times, extra_http_headers=extra_http_headers, + cdp_connect_headers=cdp_connect_headers, run_with=run_with, ai_fallback=ai_fallback, cache_key=cache_key, @@ -4310,6 +4325,7 @@ class WorkflowService: parent_workflow_run_id=parent_workflow_run_id, max_screenshot_scrolling_times=workflow_request.max_screenshot_scrolls, extra_http_headers=workflow_request.extra_http_headers, + cdp_connect_headers=workflow_request.cdp_connect_headers, browser_address=workflow_request.browser_address, sequential_key=sequential_key, run_with=workflow_request.run_with, @@ -5223,6 +5239,7 @@ class WorkflowService: totp_verification_url=workflow_run.totp_verification_url, totp_identifier=workflow_run.totp_identifier, extra_http_headers=workflow_run.extra_http_headers, + cdp_connect_headers=workflow_run.cdp_connect_headers, queued_at=workflow_run.queued_at, started_at=workflow_run.started_at, finished_at=workflow_run.finished_at, @@ -5778,6 +5795,17 @@ class WorkflowService: if existing_latest_workflow: existing_version = existing_latest_workflow.version + # Missing field inherits the stored dict; an explicit dict (possibly + # with mask sentinels for unedited keys) is resolved entry-by-entry + # against the stored value so newly-added keys aren't dropped. + if request.cdp_connect_headers is None: + effective_cdp_connect_headers = existing_latest_workflow.cdp_connect_headers + else: + effective_cdp_connect_headers = merge_masked_headers( + request.cdp_connect_headers, + existing_latest_workflow.cdp_connect_headers, + ) + # NOTE: it's only potential, as it may be immediately deleted! potential_workflow = await self.create_workflow( title=title, @@ -5793,6 +5821,7 @@ class WorkflowService: model=request.model, max_screenshot_scrolling_times=request.max_screenshot_scrolls, extra_http_headers=request.extra_http_headers, + cdp_connect_headers=effective_cdp_connect_headers, workflow_permanent_id=existing_latest_workflow.workflow_permanent_id, version=existing_version + 1, is_saved_task=request.is_saved_task, @@ -5812,6 +5841,10 @@ class WorkflowService: edited_by=edited_by, ) else: + # No existing workflow to inherit from; merge_masked_headers drops + # any keys whose value is the mask sentinel so we never persist the + # literal "***" placeholder from a misbehaving client. + new_cdp_connect_headers = merge_masked_headers(request.cdp_connect_headers, None) # NOTE: it's only potential, as it may be immediately deleted! potential_workflow = await self.create_workflow( title=title, @@ -5827,6 +5860,7 @@ class WorkflowService: model=request.model, max_screenshot_scrolling_times=request.max_screenshot_scrolls, extra_http_headers=request.extra_http_headers, + cdp_connect_headers=new_cdp_connect_headers, is_saved_task=request.is_saved_task, status=request.status, run_with=request.run_with, @@ -5904,6 +5938,7 @@ class WorkflowService: proxy_location: ProxyLocationInput = None, max_screenshot_scrolling_times: int | None = None, extra_http_headers: dict[str, str] | None = None, + cdp_connect_headers: dict[str, str] | None = None, run_with: str | None = None, status: WorkflowStatus = WorkflowStatus.published, ) -> Workflow: @@ -5921,6 +5956,7 @@ class WorkflowService: status=status, max_screenshot_scrolls=max_screenshot_scrolling_times, extra_http_headers=extra_http_headers, + cdp_connect_headers=cdp_connect_headers, run_with=run_with, ) return await app.WORKFLOW_SERVICE.create_workflow_from_request( diff --git a/skyvern/schemas/run_blocks.py b/skyvern/schemas/run_blocks.py index f7b4baaeb..b7359850e 100644 --- a/skyvern/schemas/run_blocks.py +++ b/skyvern/schemas/run_blocks.py @@ -31,6 +31,13 @@ class BaseRunBlockRequest(BaseModel): extra_http_headers: dict[str, str] | None = Field( default=None, description="Additional HTTP headers to include in requests" ) + cdp_connect_headers: dict[str, str] | None = Field( + default=None, + description=( + "HTTP headers attached ONLY to the CDP WebSocket handshake when connecting " + "to a remote browser via browser_address. Never forwarded to target websites." + ), + ) max_screenshot_scrolling_times: int | None = Field( default=None, description="Maximum number of times to scroll for screenshots" ) diff --git a/skyvern/schemas/runs.py b/skyvern/schemas/runs.py index 54860de83..d18d69fac 100644 --- a/skyvern/schemas/runs.py +++ b/skyvern/schemas/runs.py @@ -3,7 +3,7 @@ from __future__ import annotations from datetime import datetime from typing import TYPE_CHECKING, Annotated, Any, Literal, TypeAlias, Union -from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator +from pydantic import BaseModel, ConfigDict, Field, field_serializer, field_validator, model_validator from skyvern.forge.sdk.schemas.files import FileInfo from skyvern.forge.sdk.workflow.models.validators import normalize_run_metadata, normalize_run_with @@ -45,6 +45,7 @@ from skyvern.schemas.run_enums import ( # noqa: F401 RunStatus, RunType, ) +from skyvern.utils.secret_headers import mask_header_values from skyvern.utils.url_validators import validate_url # Type checkers need string Literal values, while pydantic's discriminated @@ -128,6 +129,15 @@ class TaskRunRequest(BaseModel): default=None, description="The extra HTTP headers for the requests in browser.", ) + cdp_connect_headers: dict[str, str] | None = Field( + default=None, + description=( + "HTTP headers attached ONLY to the CDP WebSocket handshake when connecting to " + "a remote browser via browser_address. Use this for browser-provider auth " + "(e.g., x-api-key for Skyvern Cloud, Browserless, or similar). These headers " + "are NEVER forwarded to target websites." + ), + ) publish_workflow: bool = Field( default=False, description="Whether to publish this task as a reusable workflow. Only available for skyvern-2.0.", @@ -174,6 +184,10 @@ class TaskRunRequest(BaseModel): return validate_url(url) + @field_serializer("cdp_connect_headers") + def _mask_cdp_connect_headers(self, headers: dict[str, str] | None) -> dict[str, str] | None: + return mask_header_values(headers) + @model_validator(mode="after") def _force_v2_for_publish_workflow(self) -> TaskRunRequest: if self.publish_workflow and self.engine != RunEngine.skyvern_v2: @@ -222,6 +236,15 @@ class WorkflowRunRequest(BaseModel): default=None, description="The extra HTTP headers for the requests in browser.", ) + cdp_connect_headers: dict[str, str] | None = Field( + default=None, + description=( + "HTTP headers attached ONLY to the CDP WebSocket handshake when connecting to " + "a remote browser via browser_address. Use this for browser-provider auth " + "(e.g., x-api-key for Skyvern Cloud, Browserless, or similar). These headers " + "are NEVER forwarded to target websites." + ), + ) browser_address: str | None = Field( default=None, description="The CDP address for the workflow run.", @@ -260,6 +283,10 @@ class WorkflowRunRequest(BaseModel): return url return validate_url(url) + @field_serializer("cdp_connect_headers") + def _mask_cdp_connect_headers(self, headers: dict[str, str] | None) -> dict[str, str] | None: + return mask_header_values(headers) + class BlockRunRequest(WorkflowRunRequest): block_labels: list[str] = Field( diff --git a/skyvern/schemas/workflows.py b/skyvern/schemas/workflows.py index 7f3e30324..0877cb2f2 100644 --- a/skyvern/schemas/workflows.py +++ b/skyvern/schemas/workflows.py @@ -5,7 +5,7 @@ from enum import StrEnum from typing import Annotated, Any, Literal import structlog -from pydantic import BaseModel, Field, field_validator, model_validator +from pydantic import BaseModel, Field, field_serializer, field_validator, model_validator from skyvern.config import settings from skyvern.forge.sdk.api.llm.config_registry import LLMConfigRegistry @@ -13,6 +13,7 @@ from skyvern.forge.sdk.settings_manager import SettingsManager from skyvern.forge.sdk.workflow.models.parameter import OutputParameter, ParameterType, WorkflowParameterType from skyvern.forge.sdk.workflow.models.validators import normalize_run_with from skyvern.schemas.runs import GeoTarget, ProxyLocation, RunEngine +from skyvern.utils.secret_headers import mask_header_values from skyvern.utils.strings import sanitize_identifier from skyvern.utils.templating import replace_jinja_reference @@ -1174,6 +1175,7 @@ class WorkflowCreateYAMLRequest(BaseModel): is_saved_task: bool = False max_screenshot_scrolls: int | None = None extra_http_headers: dict[str, str] | None = None + cdp_connect_headers: dict[str, str] | None = None status: WorkflowStatus = WorkflowStatus.published run_with: str = "agent" ai_fallback: bool = True @@ -1190,6 +1192,10 @@ class WorkflowCreateYAMLRequest(BaseModel): def _normalize_run_with(cls, v: str | None) -> str: return normalize_run_with(v) + @field_serializer("cdp_connect_headers") + def _mask_cdp_connect_headers(self, headers: dict[str, str] | None) -> dict[str, str] | None: + return mask_header_values(headers) + class WorkflowRequest(BaseModel): json_definition: WorkflowCreateYAMLRequest | None = Field( diff --git a/skyvern/services/task_v2_service.py b/skyvern/services/task_v2_service.py index 4fe2124fa..ec9fa2e9b 100644 --- a/skyvern/services/task_v2_service.py +++ b/skyvern/services/task_v2_service.py @@ -260,6 +260,7 @@ async def initialize_task_v2( max_screenshot_scrolling_times: int | None = None, browser_session_id: str | None = None, extra_http_headers: dict[str, str] | None = None, + cdp_connect_headers: dict[str, str] | None = None, browser_address: str | None = None, run_with: str | None = None, trigger_type: WorkflowRunTriggerType | None = None, @@ -278,6 +279,7 @@ async def initialize_task_v2( model=model, max_screenshot_scrolling_times=max_screenshot_scrolling_times, extra_http_headers=extra_http_headers, + cdp_connect_headers=cdp_connect_headers, browser_address=browser_address, run_with=run_with, ) @@ -299,6 +301,7 @@ async def initialize_task_v2( status=workflow_status, max_screenshot_scrolling_times=max_screenshot_scrolling_times, extra_http_headers=extra_http_headers, + cdp_connect_headers=cdp_connect_headers, run_with=run_with, ) workflow_run = await app.WORKFLOW_SERVICE.setup_workflow_run( @@ -307,6 +310,7 @@ async def initialize_task_v2( max_screenshot_scrolls=max_screenshot_scrolling_times, browser_session_id=browser_session_id, extra_http_headers=extra_http_headers, + cdp_connect_headers=cdp_connect_headers, browser_address=browser_address, run_with=run_with, ), diff --git a/skyvern/utils/secret_headers.py b/skyvern/utils/secret_headers.py new file mode 100644 index 000000000..3981d2f4c --- /dev/null +++ b/skyvern/utils/secret_headers.py @@ -0,0 +1,39 @@ +SECRET_HEADER_MASK = "***" + + +def mask_header_values(headers: dict[str, str] | None) -> dict[str, str] | None: + """Replace every value in a header dict with ``SECRET_HEADER_MASK``. + + Used by Pydantic ``field_serializer`` hooks on ``cdp_connect_headers`` so that + API responses never echo browser-provider credentials (e.g. ``x-api-key``) + back to clients polling status endpoints. The serializer fires on + ``model_dump``/JSON serialization only, so internal attribute access still + returns the raw value used for the CDP handshake. + """ + if not headers: + return headers + return {k: SECRET_HEADER_MASK for k in headers} + + +def merge_masked_headers( + new_headers: dict[str, str] | None, + stored_headers: dict[str, str] | None, +) -> dict[str, str] | None: + """Resolve ``SECRET_HEADER_MASK`` entries in ``new_headers`` against ``stored_headers`` per-key. + + Masked entries fall back to the stored value (preserving credentials on round-trip) + or are dropped when no stored value exists; non-masked entries pass through verbatim. + Returns ``None`` only when ``new_headers`` is ``None`` so callers can distinguish + "field omitted" from an explicit empty/cleared dict. + """ + if new_headers is None: + return None + stored = stored_headers or {} + merged: dict[str, str] = {} + for key, value in new_headers.items(): + if value == SECRET_HEADER_MASK: + if key in stored: + merged[key] = stored[key] + continue + merged[key] = value + return merged diff --git a/skyvern/webeye/browser_factory.py b/skyvern/webeye/browser_factory.py index ac07bd687..a07c80334 100644 --- a/skyvern/webeye/browser_factory.py +++ b/skyvern/webeye/browser_factory.py @@ -550,6 +550,7 @@ async def _create_headless_chromium( playwright: Playwright, proxy_location: ProxyLocation | None = None, extra_http_headers: dict[str, str] | None = None, + cdp_connect_headers: dict[str, str] | None = None, **kwargs: dict, ) -> tuple[BrowserContext, BrowserArtifacts, BrowserCleanupFunc]: if browser_address := kwargs.get("browser_address"): @@ -557,6 +558,7 @@ async def _create_headless_chromium( playwright, remote_browser_url=str(browser_address), extra_http_headers=extra_http_headers, + cdp_connect_headers=cdp_connect_headers, apply_download_behaviour=True, ) @@ -639,6 +641,7 @@ async def _create_headful_chromium( playwright: Playwright, proxy_location: ProxyLocation | None = None, extra_http_headers: dict[str, str] | None = None, + cdp_connect_headers: dict[str, str] | None = None, **kwargs: dict, ) -> tuple[BrowserContext, BrowserArtifacts, BrowserCleanupFunc]: if browser_address := kwargs.get("browser_address"): @@ -646,6 +649,7 @@ async def _create_headful_chromium( playwright, remote_browser_url=str(browser_address), extra_http_headers=extra_http_headers, + cdp_connect_headers=cdp_connect_headers, apply_download_behaviour=True, ) @@ -756,6 +760,7 @@ async def _create_cdp_connection_browser( playwright: Playwright, proxy_location: ProxyLocation | None = None, extra_http_headers: dict[str, str] | None = None, + cdp_connect_headers: dict[str, str] | None = None, **kwargs: dict, ) -> tuple[BrowserContext, BrowserArtifacts, BrowserCleanupFunc]: if browser_address := kwargs.get("browser_address"): @@ -763,6 +768,7 @@ async def _create_cdp_connection_browser( playwright, remote_browser_url=str(browser_address), extra_http_headers=extra_http_headers, + cdp_connect_headers=cdp_connect_headers, apply_download_behaviour=True, ) @@ -813,13 +819,19 @@ async def _create_cdp_connection_browser( else: LOG.info("Port 9222 is in use, using existing browser") - return await _connect_to_cdp_browser(playwright, settings.BROWSER_REMOTE_DEBUGGING_URL, extra_http_headers) + return await _connect_to_cdp_browser( + playwright, + settings.BROWSER_REMOTE_DEBUGGING_URL, + extra_http_headers=extra_http_headers, + cdp_connect_headers=cdp_connect_headers, + ) async def _connect_to_cdp_browser( playwright: Playwright, remote_browser_url: str, extra_http_headers: dict[str, str] | None = None, + cdp_connect_headers: dict[str, str] | None = None, apply_download_behaviour: bool = False, ) -> tuple[BrowserContext, BrowserArtifacts, BrowserCleanupFunc]: parsed_headers = parse_extra_headers(extra_http_headers) @@ -833,10 +845,20 @@ async def _connect_to_cdp_browser( ) LOG.info("Connecting browser CDP connection", remote_browser_url=remote_browser_url) + cdp_headers = build_cdp_connect_headers(settings.BROWSER_REMOTE_DEBUGGING_HOST_HEADER) or {} + if cdp_connect_headers: + # HTTP header names are case-insensitive, so a plain ``{**user, **server}`` + # would still leak a user-supplied ``host`` alongside our managed ``Host``. + # Drop any user key that collides with a server key on a lowercased compare. + reserved_keys = {key.lower() for key in cdp_headers} + filtered_user_headers = { + key: value for key, value in cdp_connect_headers.items() if key.lower() not in reserved_keys + } + cdp_headers = {**filtered_user_headers, **cdp_headers} browser = await _connect_over_cdp_with_diagnostics( playwright, remote_browser_url, - headers=build_cdp_connect_headers(settings.BROWSER_REMOTE_DEBUGGING_HOST_HEADER), + headers=cdp_headers or None, timeout_ms=settings.BROWSER_CDP_CONNECT_TIMEOUT_MS, ) diff --git a/skyvern/webeye/browser_state.py b/skyvern/webeye/browser_state.py index f16d696c5..a430e6276 100644 --- a/skyvern/webeye/browser_state.py +++ b/skyvern/webeye/browser_state.py @@ -29,6 +29,7 @@ class BrowserState(Protocol): script_id: str | None = None, organization_id: str | None = None, extra_http_headers: dict[str, str] | None = None, + cdp_connect_headers: dict[str, str] | None = None, browser_address: str | None = None, browser_profile_id: str | None = None, ) -> None: ... @@ -57,6 +58,7 @@ class BrowserState(Protocol): script_id: str | None = None, organization_id: str | None = None, extra_http_headers: dict[str, str] | None = None, + cdp_connect_headers: dict[str, str] | None = None, browser_address: str | None = None, browser_profile_id: str | None = None, ) -> Page: ... diff --git a/skyvern/webeye/real_browser_manager.py b/skyvern/webeye/real_browser_manager.py index 15264654b..4e2369455 100644 --- a/skyvern/webeye/real_browser_manager.py +++ b/skyvern/webeye/real_browser_manager.py @@ -34,6 +34,7 @@ class RealBrowserManager(BrowserManager): script_id: str | None = None, organization_id: str | None = None, extra_http_headers: dict[str, str] | None = None, + cdp_connect_headers: dict[str, str] | None = None, browser_address: str | None = None, browser_profile_id: str | None = None, allow_content_blocking_extensions: bool = True, @@ -53,6 +54,7 @@ class RealBrowserManager(BrowserManager): script_id=script_id, organization_id=organization_id, extra_http_headers=extra_http_headers, + cdp_connect_headers=cdp_connect_headers, browser_address=browser_address, browser_profile_id=browser_profile_id, allow_content_blocking_extensions=allow_content_blocking_extensions, @@ -140,6 +142,7 @@ class RealBrowserManager(BrowserManager): workflow_permanent_id=task.workflow_permanent_id, organization_id=task.organization_id, extra_http_headers=task.extra_http_headers, + cdp_connect_headers=task.cdp_connect_headers, browser_address=task.browser_address, ) @@ -162,6 +165,7 @@ class RealBrowserManager(BrowserManager): workflow_permanent_id=task.workflow_permanent_id, organization_id=task.organization_id, extra_http_headers=task.extra_http_headers, + cdp_connect_headers=task.cdp_connect_headers, browser_address=task.browser_address, ) return browser_state @@ -298,6 +302,7 @@ class RealBrowserManager(BrowserManager): workflow_permanent_id=workflow_run.workflow_permanent_id, organization_id=workflow_run.organization_id, extra_http_headers=workflow_run.extra_http_headers, + cdp_connect_headers=workflow_run.cdp_connect_headers, browser_address=workflow_run.browser_address, browser_profile_id=browser_profile_id, allow_content_blocking_extensions=effective_allow_content_blocking_extensions, @@ -326,6 +331,7 @@ class RealBrowserManager(BrowserManager): workflow_permanent_id=workflow_run.workflow_permanent_id, organization_id=workflow_run.organization_id, extra_http_headers=workflow_run.extra_http_headers, + cdp_connect_headers=workflow_run.cdp_connect_headers, browser_address=workflow_run.browser_address, browser_profile_id=browser_profile_id, ) diff --git a/skyvern/webeye/real_browser_state.py b/skyvern/webeye/real_browser_state.py index e0d9ac9a4..4f0113ea8 100644 --- a/skyvern/webeye/real_browser_state.py +++ b/skyvern/webeye/real_browser_state.py @@ -85,6 +85,7 @@ class RealBrowserState(BrowserState): script_id: str | None = None, organization_id: str | None = None, extra_http_headers: dict[str, str] | None = None, + cdp_connect_headers: dict[str, str] | None = None, browser_address: str | None = None, browser_profile_id: str | None = None, ) -> None: @@ -104,6 +105,7 @@ class RealBrowserState(BrowserState): script_id=script_id, organization_id=organization_id, extra_http_headers=extra_http_headers, + cdp_connect_headers=cdp_connect_headers, browser_address=browser_address, browser_profile_id=browser_profile_id, allow_content_blocking_extensions=self.allow_content_blocking_extensions, @@ -255,6 +257,7 @@ class RealBrowserState(BrowserState): script_id: str | None = None, organization_id: str | None = None, extra_http_headers: dict[str, str] | None = None, + cdp_connect_headers: dict[str, str] | None = None, browser_address: str | None = None, browser_profile_id: str | None = None, ) -> Page: @@ -272,6 +275,7 @@ class RealBrowserState(BrowserState): script_id=script_id, organization_id=organization_id, extra_http_headers=extra_http_headers, + cdp_connect_headers=cdp_connect_headers, browser_address=browser_address, browser_profile_id=browser_profile_id, ) @@ -293,6 +297,7 @@ class RealBrowserState(BrowserState): script_id=script_id, organization_id=organization_id, extra_http_headers=extra_http_headers, + cdp_connect_headers=cdp_connect_headers, browser_address=browser_address, browser_profile_id=browser_profile_id, ) @@ -311,6 +316,7 @@ class RealBrowserState(BrowserState): script_id=script_id, organization_id=organization_id, extra_http_headers=extra_http_headers, + cdp_connect_headers=cdp_connect_headers, browser_address=browser_address, browser_profile_id=browser_profile_id, ) diff --git a/tests/unit/test_cdp_connect_headers_migration.py b/tests/unit/test_cdp_connect_headers_migration.py new file mode 100644 index 000000000..991972f74 --- /dev/null +++ b/tests/unit/test_cdp_connect_headers_migration.py @@ -0,0 +1,170 @@ +import importlib.util +from pathlib import Path + +import pytest +from sqlalchemy.exc import DBAPIError + +MIGRATION_DIR = Path(__file__).resolve().parents[2] / "alembic/versions" + + +def _find_migration_path() -> Path: + paths = [ + path + for path in sorted(MIGRATION_DIR.glob("*_add_cdp_connect_headers.py")) + if "_execute_with_retry" in path.read_text() + ] + assert len(paths) == 1 + return paths[0] + + +class _OrigLockError(Exception): + sqlstate = "55P03" + + +class _OrigPgcodeLockError(Exception): + pgcode = "55P03" + + +class _FakeOp: + def __init__( + self, lock_failures: int = 0, fail_rollback: bool = False, lock_error: Exception | None = None + ) -> None: + self.lock_failures = lock_failures + self.fail_rollback = fail_rollback + self.lock_error = lock_error or _OrigLockError() + self.statements: list[str] = [] + + def execute(self, statement: str) -> None: + self.statements.append(statement) + if statement == "ROLLBACK" and self.fail_rollback: + raise RuntimeError("rollback failed") + if statement.startswith("LOCK TABLE") and self.lock_failures: + self.lock_failures -= 1 + raise DBAPIError(statement, None, self.lock_error) + + +def _load_migration_module(): + spec = importlib.util.spec_from_file_location("cdp_connect_headers_migration", _find_migration_path()) + assert spec is not None + assert spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def test_retry_uses_nowait_lock_and_rolls_back_failed_attempts(monkeypatch: pytest.MonkeyPatch) -> None: + migration = _load_migration_module() + fake_op = _FakeOp(lock_failures=2) + sleeps: list[int] = [] + + monkeypatch.setattr(migration, "op", fake_op) + monkeypatch.setattr(migration.time, "sleep", sleeps.append) + monkeypatch.setattr(migration.time, "monotonic", lambda: 0.0) + + migration._execute_with_retry( + "tasks", + 'ALTER TABLE "tasks" ADD COLUMN IF NOT EXISTS cdp_connect_headers JSON', + deadline=100.0, + ) + + assert sleeps == [5, 5] + assert fake_op.statements == [ + "BEGIN", + "SET LOCAL statement_timeout = '5s'", + 'LOCK TABLE "tasks" IN ACCESS EXCLUSIVE MODE NOWAIT', + "ROLLBACK", + "BEGIN", + "SET LOCAL statement_timeout = '5s'", + 'LOCK TABLE "tasks" IN ACCESS EXCLUSIVE MODE NOWAIT', + "ROLLBACK", + "BEGIN", + "SET LOCAL statement_timeout = '5s'", + 'LOCK TABLE "tasks" IN ACCESS EXCLUSIVE MODE NOWAIT', + 'ALTER TABLE "tasks" ADD COLUMN IF NOT EXISTS cdp_connect_headers JSON', + "COMMIT", + ] + + +def test_retry_stops_when_deadline_is_exhausted(monkeypatch: pytest.MonkeyPatch) -> None: + migration = _load_migration_module() + fake_op = _FakeOp(lock_failures=1) + sleeps: list[int] = [] + + monkeypatch.setattr(migration, "op", fake_op) + monkeypatch.setattr(migration.time, "sleep", sleeps.append) + monkeypatch.setattr(migration.time, "monotonic", lambda: 101.0) + + with pytest.raises(DBAPIError): + migration._execute_with_retry( + "observer_cruises", + 'ALTER TABLE "observer_cruises" ADD COLUMN IF NOT EXISTS cdp_connect_headers JSON', + deadline=100.0, + ) + + assert sleeps == [] + assert fake_op.statements == [ + "BEGIN", + "SET LOCAL statement_timeout = '5s'", + 'LOCK TABLE "observer_cruises" IN ACCESS EXCLUSIVE MODE NOWAIT', + "ROLLBACK", + ] + + +def test_retry_recognizes_pgcode_lock_errors(monkeypatch: pytest.MonkeyPatch) -> None: + migration = _load_migration_module() + fake_op = _FakeOp(lock_failures=1, lock_error=_OrigPgcodeLockError()) + sleeps: list[int] = [] + + monkeypatch.setattr(migration, "op", fake_op) + monkeypatch.setattr(migration.time, "sleep", sleeps.append) + monkeypatch.setattr(migration.time, "monotonic", lambda: 0.0) + + migration._execute_with_retry( + "tasks", + 'ALTER TABLE "tasks" ADD COLUMN IF NOT EXISTS cdp_connect_headers JSON', + deadline=100.0, + ) + + assert sleeps == [5] + assert fake_op.statements == [ + "BEGIN", + "SET LOCAL statement_timeout = '5s'", + 'LOCK TABLE "tasks" IN ACCESS EXCLUSIVE MODE NOWAIT', + "ROLLBACK", + "BEGIN", + "SET LOCAL statement_timeout = '5s'", + 'LOCK TABLE "tasks" IN ACCESS EXCLUSIVE MODE NOWAIT', + 'ALTER TABLE "tasks" ADD COLUMN IF NOT EXISTS cdp_connect_headers JSON', + "COMMIT", + ] + + +def test_retry_aborts_when_rollback_fails(monkeypatch: pytest.MonkeyPatch) -> None: + migration = _load_migration_module() + fake_op = _FakeOp(lock_failures=1, fail_rollback=True) + sleeps: list[int] = [] + + monkeypatch.setattr(migration, "op", fake_op) + monkeypatch.setattr(migration.time, "sleep", sleeps.append) + monkeypatch.setattr(migration.time, "monotonic", lambda: 0.0) + + with pytest.raises(RuntimeError, match="rollback failed"): + migration._execute_with_retry( + "tasks", + 'ALTER TABLE "tasks" ADD COLUMN IF NOT EXISTS cdp_connect_headers JSON', + deadline=100.0, + ) + + assert sleeps == [] + assert fake_op.statements == [ + "BEGIN", + "SET LOCAL statement_timeout = '5s'", + 'LOCK TABLE "tasks" IN ACCESS EXCLUSIVE MODE NOWAIT', + "ROLLBACK", + ] + + +def test_retry_budget_stays_under_ten_minutes() -> None: + migration = _load_migration_module() + + assert migration._MIGRATION_RETRY_SECONDS == 10 * 60 diff --git a/tests/unit/test_cdp_connection.py b/tests/unit/test_cdp_connection.py index 85fede7c7..6c114524c 100644 --- a/tests/unit/test_cdp_connection.py +++ b/tests/unit/test_cdp_connection.py @@ -6,6 +6,7 @@ from typing import Any, cast import pytest from playwright.async_api import Playwright +import skyvern.webeye.browser_factory as browser_factory import skyvern.webeye.cdp_connection as cdp_connection from skyvern.webeye.cdp_connection import ( build_cdp_connect_headers, @@ -107,6 +108,37 @@ async def test_connect_over_cdp_retries_resolved_host_with_headers(monkeypatch: ] +@pytest.mark.asyncio +async def test_create_cdp_connection_browser_passes_headers_to_configured_cdp_url( + monkeypatch: pytest.MonkeyPatch, +) -> None: + calls: list[tuple[str, dict[str, str] | None, dict[str, str] | None]] = [] + + async def fake_connect_to_cdp_browser( + playwright: Playwright, + remote_browser_url: str, + extra_http_headers: dict[str, str] | None = None, + cdp_connect_headers: dict[str, str] | None = None, + apply_download_behaviour: bool = False, + ) -> tuple[object, object, object]: + calls.append((remote_browser_url, extra_http_headers, cdp_connect_headers)) + return object(), object(), object() + + monkeypatch.setattr(browser_factory.settings, "BROWSER_TYPE", "chromium-headful") + monkeypatch.setattr(browser_factory.settings, "BROWSER_REMOTE_DEBUGGING_URL", "http://browser.example:9222") + monkeypatch.setattr(browser_factory, "_connect_to_cdp_browser", fake_connect_to_cdp_browser) + + await browser_factory._create_cdp_connection_browser( + cast(Playwright, object()), + extra_http_headers={"user-agent": "test"}, + cdp_connect_headers={"x-api-key": "secret"}, + ) + + assert calls == [ + ("http://browser.example:9222", {"user-agent": "test"}, {"x-api-key": "secret"}), + ] + + @pytest.mark.asyncio async def test_connect_over_cdp_accepts_direct_websocket_with_default_timeout() -> None: expected_browser = object() diff --git a/tests/unit/test_workflow_browser_profile_id.py b/tests/unit/test_workflow_browser_profile_id.py index a77291691..b592f8fb7 100644 --- a/tests/unit/test_workflow_browser_profile_id.py +++ b/tests/unit/test_workflow_browser_profile_id.py @@ -55,6 +55,20 @@ def test_workflow_create_yaml_request_accepts_browser_profile_id() -> None: assert request.browser_profile_id == "bp_abc123" +def test_workflow_create_yaml_request_masks_cdp_connect_headers_on_dump() -> None: + request = WorkflowCreateYAMLRequest( + title="test", + workflow_definition=WorkflowDefinitionYAML(parameters=[], blocks=[]), + cdp_connect_headers={"x-api-key": "secret", "authorization": "Bearer secret"}, + ) + + assert request.cdp_connect_headers == {"x-api-key": "secret", "authorization": "Bearer secret"} + assert request.model_dump()["cdp_connect_headers"] == { + "x-api-key": "***", + "authorization": "***", + } + + def _make_setup_service(workflow: SimpleNamespace) -> tuple[WorkflowService, SimpleNamespace, SimpleNamespace]: service = WorkflowService() workflow_run = SimpleNamespace(workflow_run_id="wr_test", workflow_permanent_id="wpid_test") @@ -77,6 +91,7 @@ def _make_workflow_stub(browser_profile_id: str | None) -> SimpleNamespace: proxy_location=None, webhook_callback_url=None, extra_http_headers=None, + cdp_connect_headers=None, browser_profile_id=browser_profile_id, run_with="agent", code_version=None, diff --git a/tests/unit/test_workflow_run_parameter_batching.py b/tests/unit/test_workflow_run_parameter_batching.py index b2fe3868f..5001c2a6a 100644 --- a/tests/unit/test_workflow_run_parameter_batching.py +++ b/tests/unit/test_workflow_run_parameter_batching.py @@ -56,6 +56,7 @@ def _make_service_with_mocks( proxy_location=None, webhook_callback_url=None, extra_http_headers=None, + cdp_connect_headers=None, browser_profile_id=None, run_with="agent", code_version=None,