vikarti.anatra/Skyvern
1
0
Fork 0
mirror of https://github.com/Skyvern-AI/skyvern.git synced 2026-07-09 16:09:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 5

feat(SKY-8207): consistent egress IP for Save & Reuse Session runs (#7169)

Browse source 
Co-authored-by: Suchintan Singh <suchintan@skyvern.com>
Co-authored-by: AronPerez <aperez0295@gmail.com>
This commit is contained in:
Suchintan 2026-07-07 17:45:11 -04:00 • committed by GitHub
parent 147eea686f
commit d0094d3568
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
55 changed files with 2745 additions and 146 deletions
Download patch file Download diff file Expand all files Collapse all files

12
tests/unit/forge/sdk/db/test_update_workflow_definition_fields.py
View file

@ -1,6 +1,6 @@
"""Regression tests for ``WorkflowsRepository.update_workflow_and_reconcile_definition_params``.
Covers the 17 workflow-level fields that were newly threaded through
Covers the 18 workflow-level fields that were newly threaded through
``WorkflowService.update_workflow_definition`` in the copilot-v2 stack:
- ``_UNSET``-guarded (12): ``proxy_location``, ``webhook_callback_url``,
@ -9,8 +9,8 @@ Covers the 17 workflow-level fields that were newly threaded through
``browser_profile_key``, ``totp_identifier``, ``adaptive_caching``,
``code_version``. Omitting the kwarg must leave the persisted value
unchanged; passing explicit ``None`` clears nullable columns.
- bare-``None`` (5): ``persist_browser_session``, ``run_with``,
``ai_fallback``, ``cache_key``, ``run_sequentially``. Both omitting
- bare-``None`` (6): ``persist_browser_session``, ``pin_saved_session_ip``,
``run_with``, ``ai_fallback``, ``cache_key``, ``run_sequentially``. Both omitting
the kwarg and passing ``None`` must leave the persisted value
unchanged (matches the existing ``update_workflow`` semantics).
@ -73,6 +73,7 @@ async def seeded_workflow(agent_db: AgentDB) -> dict[str, str]:
max_screenshot_scrolling_times=7,
extra_http_headers={"X-Seed": "yes"},
persist_browser_session=True,
pin_saved_session_ip=True,
browser_profile_id="bp_seed_profile",
browser_profile_key="{{ credential_id }}",
model={"model_name": "seed-model"},
@ -228,6 +229,7 @@ async def test_omitting_all_workflow_level_fields_preserves_seed(
assert workflow.max_screenshot_scrolls == 7
assert workflow.extra_http_headers == {"X-Seed": "yes"}
assert workflow.persist_browser_session is True
assert workflow.pin_saved_session_ip is True
assert workflow.browser_profile_id == "bp_seed_profile"
assert workflow.browser_profile_key == "{{ credential_id }}"
assert workflow.model == {"model_name": "seed-model"}
@ -252,6 +254,7 @@ async def test_passing_none_to_bare_none_fields_does_not_clobber(
workflow_id=seeded_workflow["workflow_id"],
organization_id=seeded_workflow["organization_id"],
persist_browser_session=None,
pin_saved_session_ip=None,
run_with=None,
ai_fallback=None,
cache_key=None,
@ -260,6 +263,7 @@ async def test_passing_none_to_bare_none_fields_does_not_clobber(
workflow = await _get(agent_db, seeded_workflow)
assert workflow.persist_browser_session is True
assert workflow.pin_saved_session_ip is True
assert workflow.run_with == "agent"
assert workflow.ai_fallback is False
assert workflow.cache_key == "seed-cache-key"
@ -318,6 +322,7 @@ async def test_setting_new_values_persists_across_all_fields(
max_screenshot_scrolling_times=12,
extra_http_headers={"X-Updated": "yes"},
persist_browser_session=False,
pin_saved_session_ip=False,
browser_profile_id="bp_new_profile",
browser_profile_key="{{ account_id }}",
model={"model_name": "new-model"},
@ -338,6 +343,7 @@ async def test_setting_new_values_persists_across_all_fields(
assert workflow.max_screenshot_scrolls == 12
assert workflow.extra_http_headers == {"X-Updated": "yes"}
assert workflow.persist_browser_session is False
assert workflow.pin_saved_session_ip is False
assert workflow.browser_profile_id == "bp_new_profile"
assert workflow.browser_profile_key == "{{ account_id }}"
assert workflow.model == {"model_name": "new-model"}

208
tests/unit/forge/sdk/db/test_workflow_runs_repository.py
View file

@ -7,10 +7,16 @@ from typing import Any
from unittest.mock import AsyncMock, MagicMock
import pytest
import pytest_asyncio
from sqlalchemy.exc import IntegrityError
from sqlalchemy.ext.asyncio import AsyncEngine
from skyvern.forge.sdk.db.agent_db import AgentDB, _build_engine
from skyvern.forge.sdk.db.models import Base, PersistentBrowserSessionModel, WorkflowRunModel
from skyvern.forge.sdk.db.repositories.workflow_runs import WorkflowRunsRepository
from skyvern.forge.sdk.schemas.persistent_browser_sessions import FORCED_WORKFLOW_SESSION_RUNNABLE_TYPE
from skyvern.forge.sdk.workflow.models.parameter import WorkflowParameter, WorkflowParameterType
from skyvern.forge.sdk.workflow.models.workflow import WorkflowRunStatus
from skyvern.schemas.runs import MAX_SEARCH_FETCH_LIMIT
@ -72,6 +78,60 @@ def _assert_not_filtering_copilot_authored_workflows(where_clause: str) -> None:
assert "workflows.edited_by" not in where_clause
@pytest_asyncio.fixture
async def sqlite_engine() -> AsyncEngine:
engine = _build_engine("sqlite+aiosqlite:///:memory:")
async with engine.begin() as conn:
await conn.run_sync(Base.metadata.create_all)
try:
yield engine
finally:
await engine.dispose()
@pytest_asyncio.fixture
async def sqlite_db(sqlite_engine: AsyncEngine) -> AgentDB:
return AgentDB("sqlite+aiosqlite:///:memory:", db_engine=sqlite_engine)
def _workflow_run_model(
*,
workflow_run_id: str,
queued_at: datetime,
browser_session_id: str | None = None,
sequential_key: str | None = None,
workflow_permanent_id: str = "wpid_test",
status: str = WorkflowRunStatus.queued.value,
) -> WorkflowRunModel:
return WorkflowRunModel(
workflow_run_id=workflow_run_id,
workflow_id="wf_test",
workflow_permanent_id=workflow_permanent_id,
organization_id="org_test",
browser_session_id=browser_session_id,
status=status,
sequential_key=sequential_key,
created_at=queued_at,
modified_at=queued_at,
queued_at=queued_at,
)
def _persistent_browser_session_model(
*,
persistent_browser_session_id: str,
runnable_type: str | None = FORCED_WORKFLOW_SESSION_RUNNABLE_TYPE,
) -> PersistentBrowserSessionModel:
now = datetime.now(tz=timezone.utc)
return PersistentBrowserSessionModel(
persistent_browser_session_id=persistent_browser_session_id,
organization_id="org_test",
runnable_type=runnable_type,
created_at=now,
modified_at=now,
)
@pytest.mark.asyncio
async def test_batch_create_uses_add_all_flush_commit_not_refresh() -> None:
"""Batch insert should use add_all + flush + commit and never call refresh."""
@ -431,6 +491,144 @@ async def test_get_workflow_runs_for_browser_session_filters_and_excludes() -> N
assert "OFFSET 5" in rendered
@pytest.mark.asyncio
async def test_get_last_queued_workflow_run_can_include_browser_session_rows(sqlite_db: AgentDB) -> None:
created_at = datetime(2026, 7, 4, 12, 0, tzinfo=timezone.utc)
async with sqlite_db.Session() as session:
session.add(
WorkflowRunModel(
workflow_run_id="wr_forced_prior",
workflow_id="wf_test",
workflow_permanent_id="wpid_test",
organization_id="org_test",
browser_session_id="pbs_forced",
status=WorkflowRunStatus.queued.value,
sequential_key="cred_a",
created_at=created_at,
modified_at=created_at,
queued_at=created_at,
)
)
await session.commit()
default_result = await sqlite_db.workflow_runs.get_last_queued_workflow_run(
"wpid_test",
"org_test",
"cred_a",
)
included_result = await sqlite_db.workflow_runs.get_last_queued_workflow_run(
"wpid_test",
"org_test",
"cred_a",
include_browser_session_rows=True,
)
assert default_result is None
assert included_result is not None
assert included_result.workflow_run_id == "wr_forced_prior"
@pytest.mark.asyncio
async def test_get_blocking_sequential_workflow_run_forced_key_lane_includes_session_rows(
sqlite_db: AgentDB,
) -> None:
prior_queued_at = datetime(2026, 7, 4, 12, 0, tzinfo=timezone.utc)
self_queued_at = datetime(2026, 7, 4, 12, 1, tzinfo=timezone.utc)
async with sqlite_db.Session() as session:
session.add_all(
[
_persistent_browser_session_model(persistent_browser_session_id="pbs_prior"),
_persistent_browser_session_model(persistent_browser_session_id="pbs_self"),
_workflow_run_model(
workflow_run_id="wr_prior",
browser_session_id="pbs_prior",
sequential_key="cred_a",
queued_at=prior_queued_at,
),
_workflow_run_model(
workflow_run_id="wr_self",
browser_session_id="pbs_self",
sequential_key="cred_a",
queued_at=self_queued_at,
),
]
)
await session.commit()
blocker = await sqlite_db.workflow_runs.get_blocking_sequential_workflow_run("wr_self")
assert blocker is not None
assert blocker.workflow_run_id == "wr_prior"
@pytest.mark.asyncio
async def test_get_blocking_sequential_workflow_run_prior_forced_visible_to_later_forced(
sqlite_db: AgentDB,
) -> None:
prior_queued_at = datetime(2026, 7, 4, 12, 0, tzinfo=timezone.utc)
self_queued_at = datetime(2026, 7, 4, 12, 1, tzinfo=timezone.utc)
async with sqlite_db.Session() as session:
session.add_all(
[
_persistent_browser_session_model(persistent_browser_session_id="pbs_forced_prior"),
_persistent_browser_session_model(persistent_browser_session_id="pbs_forced_self"),
_workflow_run_model(
workflow_run_id="wr_forced_prior",
browser_session_id="pbs_forced_prior",
sequential_key="cred_a",
queued_at=prior_queued_at,
),
_workflow_run_model(
workflow_run_id="wr_forced_self",
browser_session_id="pbs_forced_self",
sequential_key="cred_a",
queued_at=self_queued_at,
),
]
)
await session.commit()
blocker = await sqlite_db.workflow_runs.get_blocking_sequential_workflow_run("wr_forced_self")
assert blocker is not None
assert blocker.workflow_run_id == "wr_forced_prior"
@pytest.mark.asyncio
async def test_get_blocking_sequential_workflow_run_forced_whole_workflow_lane_includes_session_rows(
sqlite_db: AgentDB,
) -> None:
prior_queued_at = datetime(2026, 7, 4, 12, 0, tzinfo=timezone.utc)
self_queued_at = datetime(2026, 7, 4, 12, 1, tzinfo=timezone.utc)
async with sqlite_db.Session() as session:
session.add_all(
[
_persistent_browser_session_model(persistent_browser_session_id="pbs_prior"),
_persistent_browser_session_model(persistent_browser_session_id="pbs_self"),
_workflow_run_model(
workflow_run_id="wr_prior",
browser_session_id="pbs_prior",
queued_at=prior_queued_at,
),
_workflow_run_model(
workflow_run_id="wr_self",
browser_session_id="pbs_self",
queued_at=self_queued_at,
),
]
)
await session.commit()
blocker = await sqlite_db.workflow_runs.get_blocking_sequential_workflow_run("wr_self")
assert blocker is not None
assert blocker.workflow_run_id == "wr_prior"
@pytest.mark.asyncio
async def test_get_blocking_sequential_workflow_run_scans_earlier_active_same_key() -> None:
"""SKY-10799: the sequential gate scans ALL earlier-queued same-key runs still in flight
@ -502,10 +700,16 @@ async def test_get_blocking_sequential_workflow_run_prefers_browser_session_lane
fake_run.queued_at = datetime(2026, 6, 8, 18, 53, 50, tzinfo=timezone.utc)
calls: list[Any] = []
persistent_browser_session = MagicMock()
persistent_browser_session.runnable_type = "user_browser_session"
async def _scalars(query: Any) -> Any:
calls.append(query)
return _Result(fake_run) if len(calls) == 1 else _Result(None)
if len(calls) == 1:
return _Result(fake_run)
if len(calls) == 2:
return _Result(persistent_browser_session)
return _Result(None)
session = MagicMock()
session.scalars = AsyncMock(side_effect=_scalars)
@ -514,7 +718,7 @@ async def test_get_blocking_sequential_workflow_run_prefers_browser_session_lane
await repo.get_blocking_sequential_workflow_run("wr_self")
where_clause = _where_clause_sql(calls[1])
where_clause = _where_clause_sql(calls[2])
assert "workflow_runs.browser_session_id = 'pbs_test'" in where_clause
assert "sequential_key" not in where_clause

2
tests/unit/test_copilot_code_block_persist_seam.py
View file

@ -3126,6 +3126,7 @@ class TestCodeRepairProgressClassification:
totp_verification_url=None,
totp_identifier=None,
persist_browser_session=False,
pin_saved_session_ip=False,
browser_profile_id=None,
browser_profile_key=None,
model=None,
@ -3209,6 +3210,7 @@ class TestCodeRepairProgressClassification:
totp_verification_url=None,
totp_identifier=None,
persist_browser_session=False,
pin_saved_session_ip=False,
browser_profile_id=None,
browser_profile_key=None,
model=None,

192
tests/unit/test_credential_rotation.py
View file

@ -18,8 +18,10 @@ from skyvern.forge.sdk.db.models import Base, WorkflowRunCredentialSelectionMode
from skyvern.forge.sdk.db.repositories.workflow_run_credential_selections import (
WorkflowRunCredentialSelectionsRepository,
)
from skyvern.forge.sdk.workflow.browser_profile_key import build_browser_profile_key_digest
from skyvern.forge.sdk.workflow.credential_selection import select_credential_for_run
from skyvern.forge.sdk.workflow.models.parameter import CredentialParameter
from skyvern.forge.sdk.workflow.models.workflow import WorkflowRequestBody
from skyvern.forge.sdk.workflow.service import WorkflowService
from skyvern.forge.sdk.workflow.workflow_definition_converter import convert_workflow_definition
from skyvern.schemas.workflows import CredentialParameterYAML, WorkflowDefinitionYAML
@ -268,6 +270,45 @@ async def test_round_robin_repository_idempotent_recall_returns_existing(sqlite_
assert count == 1
@pytest.mark.asyncio
async def test_repository_get_selections_for_run_returns_mapping(sqlite_db: AgentDB) -> None:
repo = sqlite_db.workflow_run_credential_selections
async with sqlite_db.Session() as session:
session.add_all(
[
WorkflowRunCredentialSelectionModel(
organization_id="org_test",
workflow_run_id="wr_test",
workflow_permanent_id="wpid_test",
parameter_key="login_cred",
credential_id="cred_a",
),
WorkflowRunCredentialSelectionModel(
organization_id="org_test",
workflow_run_id="wr_test",
workflow_permanent_id="wpid_test",
parameter_key="backup_cred",
credential_id="cred_b",
),
WorkflowRunCredentialSelectionModel(
organization_id="org_test",
workflow_run_id="wr_other",
workflow_permanent_id="wpid_test",
parameter_key="login_cred",
credential_id="cred_other",
),
]
)
await session.commit()
assert await repo.get_selections_for_run("wr_test") == {
"backup_cred": "cred_b",
"login_cred": "cred_a",
}
assert await repo.get_selections_for_run("wr_missing") == {}
@pytest.mark.asyncio
async def test_rotation_advisory_lock_skips_non_postgres_dialect() -> None:
repo = WorkflowRunCredentialSelectionsRepository(MagicMock())
@ -421,3 +462,154 @@ async def test_resolve_login_block_credential_ids_returns_selected_rotating_id()
assert credential_ids == ["cred_b"]
context.resolve_credential_parameter_id.assert_awaited_once_with(parameter, "org_test")
def _setup_workflow_with_rotating_credential(browser_profile_key: str | None = "{{ login_cred }}") -> SimpleNamespace:
return SimpleNamespace(
workflow_id="wf_test",
workflow_permanent_id="wpid_test",
organization_id="org_test",
proxy_location=None,
webhook_callback_url=None,
extra_http_headers=None,
cdp_connect_headers=None,
browser_profile_id=None,
persist_browser_session=True,
pin_saved_session_ip=False,
browser_profile_key=browser_profile_key,
title="Workflow",
max_elapsed_time_minutes=None,
run_with="agent",
code_version=None,
adaptive_caching=False,
sequential_key=None,
workflow_definition=SimpleNamespace(parameters=[_credential_parameter(credential_ids=["cred_a", "cred_b"])]),
)
def _setup_workflow_run() -> SimpleNamespace:
return SimpleNamespace(
workflow_run_id="wr_test",
workflow_permanent_id="wpid_test",
organization_id="org_test",
browser_profile_id=None,
proxy_location=None,
)
async def _setup_rotation_profile_run(
*,
select_side_effect: str | Exception,
profile_id: str,
) -> tuple[SimpleNamespace, MagicMock]:
result, mock_app, _, caught = await _attempt_setup_rotation_profile_run(
select_side_effect=select_side_effect,
profile_id=profile_id,
)
if caught:
raise caught
assert result is not None
return result, mock_app
async def _attempt_setup_rotation_profile_run(
*,
select_side_effect: str | Exception,
profile_id: str,
browser_profile_key: str | None = "{{ login_cred }}",
) -> tuple[SimpleNamespace | None, MagicMock, WorkflowService, Exception | None]:
service = WorkflowService()
workflow = _setup_workflow_with_rotating_credential(browser_profile_key=browser_profile_key)
workflow_run = _setup_workflow_run()
updated_run_values = dict(workflow_run.__dict__)
updated_run_values["browser_profile_id"] = profile_id
updated_run = SimpleNamespace(**updated_run_values)
organization = SimpleNamespace(organization_id="org_test", organization_name="Test Org")
service.get_workflow_by_permanent_id = AsyncMock(return_value=workflow) # type: ignore[method-assign]
service.create_workflow_run = AsyncMock(return_value=workflow_run) # type: ignore[method-assign]
service.get_workflow_parameters = AsyncMock(return_value=[]) # type: ignore[method-assign]
service.create_workflow_run_parameters = AsyncMock(return_value=[]) # type: ignore[method-assign]
service.mark_workflow_run_as_failed = AsyncMock(return_value=workflow_run) # type: ignore[method-assign]
select_mock = (
AsyncMock(side_effect=select_side_effect)
if isinstance(select_side_effect, Exception)
else AsyncMock(return_value=select_side_effect)
)
with (
patch("skyvern.forge.sdk.workflow.service.app") as mock_app,
patch("skyvern.forge.sdk.workflow.service.select_credential_for_run", select_mock),
):
mock_app.EXPERIMENTATION_PROVIDER.is_feature_enabled_cached = AsyncMock(return_value=False)
mock_app.AGENT_FUNCTION.should_use_flex_llm_routing = AsyncMock(return_value=False)
mock_app.DATABASE.browser_sessions.get_or_create_managed_browser_profile = AsyncMock(
return_value=(
SimpleNamespace(browser_profile_id=profile_id, is_managed=True, proxy_session_id=None),
False,
)
)
mock_app.DATABASE.workflow_runs.update_workflow_run = AsyncMock(return_value=updated_run)
result = None
caught = None
try:
result = await service.setup_workflow_run(
request_id="req_test",
workflow_request=WorkflowRequestBody(data={}),
workflow_permanent_id="wpid_test",
organization=organization,
)
except Exception as exc:
caught = exc
return result, mock_app, service, caught
@pytest.mark.asyncio
async def test_setup_workflow_run_uses_selected_rotating_credential_for_profile_key() -> None:
result, mock_app = await _setup_rotation_profile_run(select_side_effect="cred_b", profile_id="bp_selected")
assert result.browser_profile_id == "bp_selected"
mock_app.DATABASE.browser_sessions.get_or_create_managed_browser_profile.assert_awaited_once_with(
organization_id="org_test",
workflow_permanent_id="wpid_test",
browser_profile_key_digest=build_browser_profile_key_digest("cred_b"),
name="Workflow (auto-saved: cred_b)",
)
@pytest.mark.asyncio
async def test_keyed_setup_workflow_run_fails_when_rotation_selection_fails() -> None:
result, _, service, caught = await _attempt_setup_rotation_profile_run(
select_side_effect=RuntimeError("selection failed"),
profile_id="bp_keyless",
)
assert result is None
assert isinstance(caught, RuntimeError)
assert str(caught) == "selection failed"
service.mark_workflow_run_as_failed.assert_awaited_once()
assert service.mark_workflow_run_as_failed.await_args.kwargs["workflow_run_id"] == "wr_test"
assert service.mark_workflow_run_as_failed.await_args.kwargs["failure_reason"].startswith(
"Setup workflow failed. failure reason:"
)
@pytest.mark.asyncio
async def test_keyless_setup_workflow_run_falls_back_to_keyless_profile_when_rotation_selection_fails() -> None:
result, mock_app, _, caught = await _attempt_setup_rotation_profile_run(
select_side_effect=RuntimeError("selection failed"),
profile_id="bp_keyless",
browser_profile_key=None,
)
assert caught is None
assert result is not None
assert result.browser_profile_id == "bp_keyless"
mock_app.DATABASE.browser_sessions.get_or_create_managed_browser_profile.assert_awaited_once_with(
organization_id="org_test",
workflow_permanent_id="wpid_test",
browser_profile_key_digest="",
name="Workflow (auto-saved session)",
)

16
tests/unit/test_login_block_credential_proxy_pin.py
View file

@ -35,7 +35,7 @@ def _proxy_header_hook() -> SimpleNamespace:
@pytest.mark.asyncio
async def test_login_block_credential_proxy_pin_applies_to_workflow_run() -> None:
service = WorkflowService()
workflow_run = SimpleNamespace(extra_http_headers={"X-Test": "1"}, proxy_location=None)
workflow_run = SimpleNamespace(workflow_run_id="wr_1", extra_http_headers={"X-Test": "1"}, proxy_location=None)
block = SimpleNamespace(parameters=[_credential_parameter()])
credential = SimpleNamespace(proxy_session_id="credential-pin", proxy_location=ProxyLocation.RESIDENTIAL_ISP)
@ -63,7 +63,9 @@ async def test_login_block_credential_proxy_pin_applies_to_workflow_run() -> Non
@pytest.mark.asyncio
async def test_login_block_credential_proxy_pin_respects_existing_run_header() -> None:
service = WorkflowService()
workflow_run = SimpleNamespace(extra_http_headers={"dedicated-ip": "run-pin"}, proxy_location=None)
workflow_run = SimpleNamespace(
workflow_run_id="wr_1", extra_http_headers={"dedicated-ip": "run-pin"}, proxy_location=None
)
block = SimpleNamespace(parameters=[_credential_parameter()])
with patch("skyvern.forge.sdk.workflow.service.app") as mock_app:
@ -86,7 +88,9 @@ async def test_login_block_credential_proxy_pin_respects_existing_run_header() -
@pytest.mark.asyncio
async def test_login_block_credential_proxy_pin_preserves_run_proxy_location() -> None:
service = WorkflowService()
workflow_run = SimpleNamespace(extra_http_headers=None, proxy_location=ProxyLocation.RESIDENTIAL)
workflow_run = SimpleNamespace(
workflow_run_id="wr_1", extra_http_headers=None, proxy_location=ProxyLocation.RESIDENTIAL
)
block = SimpleNamespace(parameters=[_credential_parameter()])
credential = SimpleNamespace(proxy_session_id="credential-pin", proxy_location=ProxyLocation.RESIDENTIAL_ISP)
@ -113,7 +117,7 @@ async def test_login_block_credential_proxy_pin_preserves_run_proxy_location() -
@pytest.mark.asyncio
async def test_login_block_credential_proxy_pin_prefers_linked_profile_pin() -> None:
service = WorkflowService()
workflow_run = SimpleNamespace(extra_http_headers=None, proxy_location=None)
workflow_run = SimpleNamespace(workflow_run_id="wr_1", extra_http_headers=None, proxy_location=None)
block = SimpleNamespace(parameters=[_credential_parameter()])
credential = SimpleNamespace(
proxy_session_id="credential-pin",
@ -150,7 +154,7 @@ async def test_login_block_credential_proxy_pin_prefers_linked_profile_pin() ->
@pytest.mark.asyncio
async def test_login_block_credential_proxy_pin_skips_when_linked_profile_is_unpinned() -> None:
service = WorkflowService()
workflow_run = SimpleNamespace(extra_http_headers=None, proxy_location=None)
workflow_run = SimpleNamespace(workflow_run_id="wr_1", extra_http_headers=None, proxy_location=None)
block = SimpleNamespace(parameters=[_credential_parameter()])
credential = SimpleNamespace(
proxy_session_id="credential-pin",
@ -180,7 +184,7 @@ async def test_login_block_credential_proxy_pin_skips_when_linked_profile_is_unp
@pytest.mark.asyncio
async def test_login_block_credential_proxy_pin_does_not_mutate_run_when_db_update_fails() -> None:
service = WorkflowService()
workflow_run = SimpleNamespace(extra_http_headers=None, proxy_location=None)
workflow_run = SimpleNamespace(workflow_run_id="wr_1", extra_http_headers=None, proxy_location=None)
block = SimpleNamespace(parameters=[_credential_parameter()])
credential = SimpleNamespace(proxy_session_id="credential-pin", proxy_location=ProxyLocation.RESIDENTIAL_ISP)

1111
tests/unit/test_saved_session_proxy_pin.py Normal file
View file

File diff suppressed because it is too large Load diff

49
tests/unit/test_workflow_browser_profile_id.py
View file

@ -120,6 +120,53 @@ async def test_create_workflow_from_request_preserves_enable_self_healing_when_o
assert create_workflow_mock.await_args.kwargs["enable_self_healing"] is True
@pytest.mark.asyncio
async def test_create_workflow_from_request_preserves_pin_saved_session_ip_when_omitted() -> None:
service, _ = _make_workflow_update_service(
existing_max_elapsed_time_minutes=None, existing_pin_saved_session_ip=True
)
request = WorkflowCreateYAMLRequest(
title="test",
workflow_definition=WorkflowDefinitionYAML(parameters=[], blocks=[]),
)
await service.create_workflow_from_request(
organization=cast(Any, SimpleNamespace(organization_id="org_1")),
request=request,
workflow_permanent_id="wpid_test",
)
create_workflow_mock = service.create_workflow
assert isinstance(create_workflow_mock, AsyncMock)
assert create_workflow_mock.await_args is not None
assert create_workflow_mock.await_args.kwargs["pin_saved_session_ip"] is True
@pytest.mark.asyncio
async def test_create_workflow_from_request_explicit_false_clears_pin_saved_session_ip() -> None:
service, _ = _make_workflow_update_service(
existing_max_elapsed_time_minutes=None, existing_pin_saved_session_ip=True
)
request = WorkflowCreateYAMLRequest(
title="test",
workflow_definition=WorkflowDefinitionYAML(parameters=[], blocks=[]),
pin_saved_session_ip=False,
)
await service.create_workflow_from_request(
organization=cast(Any, SimpleNamespace(organization_id="org_1")),
request=request,
workflow_permanent_id="wpid_test",
)
create_workflow_mock = service.create_workflow
assert isinstance(create_workflow_mock, AsyncMock)
assert create_workflow_mock.await_args is not None
assert create_workflow_mock.await_args.kwargs["pin_saved_session_ip"] is False
@pytest.mark.asyncio
async def test_create_workflow_from_request_explicit_false_clears_enable_self_healing() -> None:
service, _ = _make_workflow_update_service(
@ -227,6 +274,7 @@ async def test_refresh_workflow_schedule_runtime_limits_reupserts_backend_schedu
def _make_workflow_update_service(
existing_max_elapsed_time_minutes: int | None,
existing_enable_self_healing: bool = True,
existing_pin_saved_session_ip: bool = False,
) -> tuple[WorkflowService, SimpleNamespace]:
service = WorkflowService()
existing_workflow = SimpleNamespace(
@ -237,6 +285,7 @@ def _make_workflow_update_service(
code_version=None,
max_elapsed_time_minutes=existing_max_elapsed_time_minutes,
enable_self_healing=existing_enable_self_healing,
pin_saved_session_ip=existing_pin_saved_session_ip,
)
potential_workflow = SimpleNamespace(workflow_id="wf_new")
updated_workflow = SimpleNamespace(workflow_id="wf_new", workflow_permanent_id="wpid_test")

8
tests/unit/test_workflow_elapsed_timeout.py
View file

@ -188,6 +188,7 @@ async def test_execute_workflow_returns_after_elapsed_timeout_without_finally(mo
workflow = SimpleNamespace(
workflow_id="wf_1",
persist_browser_session=False,
workflow_permanent_id="wp_1",
title="Timeout workflow",
organization_id="org_1",
@ -264,6 +265,7 @@ async def test_execute_workflow_times_out_slow_pre_block_script_lookup(monkeypat
workflow = SimpleNamespace(
workflow_id="wf_1",
persist_browser_session=False,
workflow_permanent_id="wp_1",
title="Timeout workflow",
organization_id="org_1",
@ -336,6 +338,7 @@ async def test_execute_workflow_preserves_completed_status_after_post_run_timeou
workflow = SimpleNamespace(
workflow_id="wf_1",
persist_browser_session=False,
workflow_permanent_id="wp_1",
title="Timeout workflow",
organization_id="org_1",
@ -426,6 +429,7 @@ async def test_execute_workflow_preserves_timed_out_status_after_non_terminal_po
workflow = SimpleNamespace(
workflow_id="wf_1",
persist_browser_session=False,
workflow_permanent_id="wp_1",
title="Timeout workflow",
organization_id="org_1",
@ -522,6 +526,7 @@ async def test_execute_workflow_marks_timed_out_when_post_run_budget_is_exhauste
workflow = SimpleNamespace(
workflow_id="wf_1",
persist_browser_session=False,
workflow_permanent_id="wp_1",
title="Timeout workflow",
organization_id="org_1",
@ -695,6 +700,7 @@ async def test_execute_workflow_refreshes_terminal_status_after_immediate_post_r
workflow = SimpleNamespace(
workflow_id="wf_1",
persist_browser_session=False,
workflow_permanent_id="wp_1",
title="Timeout workflow",
organization_id="org_1",
@ -785,6 +791,7 @@ async def test_execute_workflow_returns_finalized_status_after_post_run_timeout(
workflow = SimpleNamespace(
workflow_id="wf_1",
persist_browser_session=False,
workflow_permanent_id="wp_1",
title="Timeout workflow",
organization_id="org_1",
@ -881,6 +888,7 @@ async def test_execute_workflow_runs_finally_for_existing_timed_out_status(
workflow = SimpleNamespace(
workflow_id="wf_1",
persist_browser_session=False,
workflow_permanent_id="wp_1",
title="Timeout workflow",
organization_id="org_1",

319
tests/unit/test_workflow_service_managed_browser_profile.py
View file

@ -1,22 +1,30 @@
from __future__ import annotations
from datetime import UTC, datetime
from types import SimpleNamespace
from unittest.mock import AsyncMock
import pytest
from skyvern.forge import app
from skyvern.forge.sdk.workflow import service as service_module
from skyvern.forge.sdk.workflow.browser_profile_key import (
build_browser_profile_key_digest,
build_workflow_browser_session_storage_key,
)
from skyvern.forge.sdk.workflow.models.block import BlockType
from skyvern.forge.sdk.workflow.service import WorkflowService
from skyvern.forge.sdk.workflow.models.workflow import WorkflowRunStatus
from skyvern.forge.sdk.workflow.service import (
CODE_BLOCK_SESSION_TIMEOUT_MINUTES,
WorkflowBrowserCleanupResult,
WorkflowService,
)
def _workflow(browser_profile_key: str | None = None) -> SimpleNamespace:
return SimpleNamespace(
persist_browser_session=True,
pin_saved_session_ip=False,
browser_profile_key=browser_profile_key,
workflow_permanent_id="wpid_test",
title="Workflow",
@ -28,6 +36,131 @@ def _workflow_run(browser_profile_id: str | None = None) -> SimpleNamespace:
workflow_run_id="wr_test",
organization_id="o_test",
browser_profile_id=browser_profile_id,
proxy_location=None,
)
def _execute_workflow() -> SimpleNamespace:
return SimpleNamespace(
workflow_id="wf_1",
persist_browser_session=True,
workflow_permanent_id="wpid_test",
title="Workflow",
organization_id="o_test",
generate_script_on_terminal=False,
model=None,
workflow_definition=SimpleNamespace(parameters=[], finally_block_label=None, blocks=[]),
)
def _execute_workflow_run(status: WorkflowRunStatus) -> SimpleNamespace:
now = datetime.now(UTC)
return SimpleNamespace(
workflow_run_id="wr_test",
workflow_id="wf_1",
workflow_permanent_id="wpid_test",
organization_id="o_test",
browser_profile_id="bp_managed",
browser_address=None,
status=status,
failure_reason=None,
ignore_inherited_workflow_system_prompt=False,
parent_workflow_run_id=None,
proxy_location=None,
max_elapsed_time_minutes=1,
started_at=now,
created_at=now,
code_gen=False,
run_with="agent",
)
def _browser_cleanup_result() -> WorkflowBrowserCleanupResult:
browser_state = SimpleNamespace(
browser_artifacts=SimpleNamespace(browser_session_dir="/tmp/fake_profile"),
browser_context=SimpleNamespace(),
)
return WorkflowBrowserCleanupResult(
browser_state=browser_state,
tasks=[],
all_workflow_task_ids=[],
child_workflow_run_ids=[],
close_browser_on_completion=True,
)
def _patch_execute_workflow_deps(
monkeypatch: pytest.MonkeyPatch,
svc: WorkflowService,
workflow: SimpleNamespace,
refreshed_run: SimpleNamespace,
) -> None:
created_run = _execute_workflow_run(WorkflowRunStatus.created)
running_run = _execute_workflow_run(WorkflowRunStatus.running)
workflow_context_manager = SimpleNamespace(
initialize_workflow_run_context=AsyncMock(),
get_workflow_run_context=lambda _workflow_run_id: SimpleNamespace(browser_session_id=None),
remove_workflow_run_context=lambda _workflow_run_id: None,
)
database = SimpleNamespace(
workflow_runs=SimpleNamespace(
get_workflow_run=AsyncMock(return_value=refreshed_run),
update_workflow_run=AsyncMock(),
),
artifacts=SimpleNamespace(claim_session_download_artifacts_for_run=AsyncMock(return_value=0)),
)
monkeypatch.setattr(service_module.app, "WORKFLOW_CONTEXT_MANAGER", workflow_context_manager)
monkeypatch.setattr(service_module.app, "DATABASE", database)
monkeypatch.setattr(service_module.app.ARTIFACT_MANAGER, "wait_for_upload_aiotasks", AsyncMock())
monkeypatch.setattr(service_module.app.STORAGE, "save_downloaded_files", AsyncMock())
monkeypatch.setattr(service_module.workflow_script_service, "workflow_has_conditionals", lambda _workflow: False)
monkeypatch.setattr(
service_module.workflow_script_service,
"get_workflow_script",
AsyncMock(return_value=(None, None, False)),
)
monkeypatch.setattr(service_module.skyvern_context, "current", lambda: None)
monkeypatch.setattr(service_module, "is_adaptive_caching", lambda _workflow, _workflow_run: False)
monkeypatch.setattr(svc, "get_workflow_run", AsyncMock(return_value=created_run))
monkeypatch.setattr(svc, "get_workflow", AsyncMock(return_value=workflow))
monkeypatch.setattr(svc, "mark_workflow_run_as_running", AsyncMock(return_value=running_run))
monkeypatch.setattr(svc, "get_workflow_run_parameter_tuples", AsyncMock(return_value=[]))
monkeypatch.setattr(svc, "get_workflow_output_parameters", AsyncMock(return_value=[]))
monkeypatch.setattr(svc, "_collect_inherited_workflow_system_prompt", AsyncMock(return_value=None))
monkeypatch.setattr(svc, "auto_create_browser_session_if_needed", AsyncMock(return_value=None))
monkeypatch.setattr(svc, "auto_create_browser_session_for_code_block_if_needed", AsyncMock(return_value=None))
monkeypatch.setattr(svc, "_browser_profile_is_managed", AsyncMock(return_value=False))
monkeypatch.setattr(svc, "_execute_workflow_blocks", AsyncMock(return_value=(refreshed_run, set())))
monkeypatch.setattr(svc, "generate_script_if_needed", AsyncMock())
monkeypatch.setattr(svc, "should_run_script", AsyncMock(return_value=False))
def _patch_browser_cleanup(monkeypatch: pytest.MonkeyPatch, svc: WorkflowService, order: list[str]) -> AsyncMock:
clean_up_browser = AsyncMock(side_effect=lambda **_kwargs: order.append("teardown") or _browser_cleanup_result())
monkeypatch.setattr(svc, "_clean_up_workflow_browser", clean_up_browser)
return clean_up_browser
def _patch_finalize(
monkeypatch: pytest.MonkeyPatch,
svc: WorkflowService,
order: list[str],
finalized_run: SimpleNamespace,
) -> None:
monkeypatch.setattr(
svc,
"_finalize_workflow_run_status",
AsyncMock(side_effect=lambda **_kwargs: order.append("finalize") or finalized_run),
)
async def _run_execute_workflow(svc: WorkflowService) -> SimpleNamespace:
return await svc.execute_workflow(
workflow_run_id="wr_test",
api_key=None,
organization=SimpleNamespace(organization_id="o_test"),
)
@ -223,6 +356,41 @@ async def test_auto_create_browser_session_for_human_interaction_loads_managed_p
timeout_minutes=61,
browser_profile_id="bp_managed",
proxy_location=None,
inherit_profile_proxy=True,
)
@pytest.mark.asyncio
async def test_auto_create_browser_session_for_code_block_loads_managed_profile_pin(
monkeypatch: pytest.MonkeyPatch,
) -> None:
create_session = AsyncMock(return_value=SimpleNamespace(persistent_browser_session_id="pbs_test"))
monkeypatch.setattr(app.PERSISTENT_SESSIONS_MANAGER, "create_session", create_session)
monkeypatch.setattr(
app.AGENT_FUNCTION,
"should_auto_create_browser_session_for_code_block",
AsyncMock(return_value=True),
)
workflow = SimpleNamespace(
workflow_id="wf_1",
workflow_permanent_id="wpid_test",
workflow_definition=SimpleNamespace(blocks=[SimpleNamespace(block_type=BlockType.CODE, loop_blocks=[])]),
)
browser_session = await WorkflowService().auto_create_browser_session_for_code_block_if_needed(
"o_test",
workflow,
workflow_run_id="wr_test",
browser_profile_id="bp_managed",
)
assert browser_session.persistent_browser_session_id == "pbs_test"
create_session.assert_awaited_once_with(
organization_id="o_test",
timeout_minutes=CODE_BLOCK_SESSION_TIMEOUT_MINUTES,
browser_profile_id="bp_managed",
proxy_location=None,
inherit_profile_proxy=True,
)
@ -251,3 +419,152 @@ async def test_browser_profile_is_managed_detects_managed_profile(monkeypatch: p
)
assert is_managed is True
@pytest.mark.asyncio
async def test_execute_workflow_persists_managed_profile_before_final_status(
monkeypatch: pytest.MonkeyPatch,
) -> None:
workflow = _execute_workflow()
completed_run = _execute_workflow_run(WorkflowRunStatus.completed)
order: list[str] = []
svc = WorkflowService()
_patch_execute_workflow_deps(monkeypatch, svc, workflow, _execute_workflow_run(WorkflowRunStatus.running))
clean_up_browser = _patch_browser_cleanup(monkeypatch, svc, order)
monkeypatch.setattr(
svc,
"_persist_workflow_browser_session_if_needed",
AsyncMock(side_effect=lambda **_kwargs: order.append("store")),
)
_patch_finalize(monkeypatch, svc, order, completed_run)
monkeypatch.setattr(svc, "persist_video_data", AsyncMock(side_effect=lambda *a, **k: order.append("video")))
monkeypatch.setattr(svc, "execute_workflow_webhook", AsyncMock(side_effect=lambda *a, **k: order.append("webhook")))
result = await _run_execute_workflow(svc)
assert result is completed_run
assert order == ["teardown", "store", "finalize", "video", "webhook"]
clean_up_browser.assert_awaited_once()
svc._persist_workflow_browser_session_if_needed.assert_awaited_once()
@pytest.mark.asyncio
async def test_execute_workflow_retries_write_back_when_pre_final_persist_fails(
monkeypatch: pytest.MonkeyPatch,
) -> None:
workflow = _execute_workflow()
completed_run = _execute_workflow_run(WorkflowRunStatus.completed)
order: list[str] = []
persist_calls = {"n": 0}
async def persist_side_effect(**_kwargs: object) -> None:
persist_calls["n"] += 1
if persist_calls["n"] == 1:
order.append("store_fail")
raise RuntimeError("storage down")
order.append("store_retry")
svc = WorkflowService()
_patch_execute_workflow_deps(monkeypatch, svc, workflow, _execute_workflow_run(WorkflowRunStatus.running))
_patch_browser_cleanup(monkeypatch, svc, order)
monkeypatch.setattr(svc, "_persist_workflow_browser_session_if_needed", AsyncMock(side_effect=persist_side_effect))
_patch_finalize(monkeypatch, svc, order, completed_run)
monkeypatch.setattr(svc, "persist_video_data", AsyncMock(side_effect=lambda *a, **k: order.append("video")))
monkeypatch.setattr(svc, "execute_workflow_webhook", AsyncMock(side_effect=lambda *a, **k: order.append("webhook")))
result = await _run_execute_workflow(svc)
assert result is completed_run
# A failed pre-final write-back must not suppress clean_up_workflow's retry after final status.
assert order == ["teardown", "store_fail", "finalize", "video", "store_retry", "webhook"]
assert svc._persist_workflow_browser_session_if_needed.await_count == 2
@pytest.mark.asyncio
async def test_execute_workflow_does_not_prestore_blob_for_failed_canceled_or_timed_out_runs(
monkeypatch: pytest.MonkeyPatch,
) -> None:
for status in (WorkflowRunStatus.failed, WorkflowRunStatus.canceled, WorkflowRunStatus.timed_out):
workflow = _execute_workflow()
terminal_run = _execute_workflow_run(status)
order: list[str] = []
svc = WorkflowService()
_patch_execute_workflow_deps(monkeypatch, svc, workflow, terminal_run)
_patch_browser_cleanup(monkeypatch, svc, order)
monkeypatch.setattr(
service_module.app.STORAGE,
"store_browser_profile",
AsyncMock(side_effect=lambda *a, **k: order.append("store_profile")),
)
monkeypatch.setattr(service_module.app.STORAGE, "store_browser_session", AsyncMock())
_patch_finalize(monkeypatch, svc, order, terminal_run)
monkeypatch.setattr(svc, "persist_video_data", AsyncMock(side_effect=lambda *a, **k: order.append("video")))
monkeypatch.setattr(svc, "execute_workflow_webhook", AsyncMock())
result = await _run_execute_workflow(svc)
assert result is terminal_run
assert order[:2] == ["finalize", "teardown"]
assert "store_profile" not in order
@pytest.mark.asyncio
async def test_execute_workflow_non_persist_workflow_tears_down_once_after_final_status(
monkeypatch: pytest.MonkeyPatch,
) -> None:
workflow = _execute_workflow()
workflow.persist_browser_session = False
completed_run = _execute_workflow_run(WorkflowRunStatus.completed)
order: list[str] = []
svc = WorkflowService()
_patch_execute_workflow_deps(monkeypatch, svc, workflow, _execute_workflow_run(WorkflowRunStatus.running))
clean_up_browser = _patch_browser_cleanup(monkeypatch, svc, order)
monkeypatch.setattr(service_module.app.STORAGE, "store_browser_profile", AsyncMock())
monkeypatch.setattr(service_module.app.STORAGE, "store_browser_session", AsyncMock())
_patch_finalize(monkeypatch, svc, order, completed_run)
monkeypatch.setattr(svc, "persist_video_data", AsyncMock(side_effect=lambda *a, **k: order.append("video")))
monkeypatch.setattr(svc, "execute_workflow_webhook", AsyncMock())
result = await _run_execute_workflow(svc)
assert result is completed_run
assert order[:2] == ["finalize", "teardown"]
clean_up_browser.assert_awaited_once()
service_module.app.STORAGE.store_browser_profile.assert_not_awaited()
service_module.app.STORAGE.store_browser_session.assert_not_awaited()
@pytest.mark.asyncio
async def test_execute_workflow_finalizes_when_pre_status_browser_cleanup_fails(
monkeypatch: pytest.MonkeyPatch,
) -> None:
workflow = _execute_workflow()
completed_run = _execute_workflow_run(WorkflowRunStatus.completed)
order: list[str] = []
def cleanup_side_effect(**_kwargs: object) -> WorkflowBrowserCleanupResult:
if "teardown_error" not in order:
order.append("teardown_error")
raise RuntimeError("cleanup failed")
order.append("teardown_cleanup")
return _browser_cleanup_result()
svc = WorkflowService()
_patch_execute_workflow_deps(monkeypatch, svc, workflow, _execute_workflow_run(WorkflowRunStatus.running))
monkeypatch.setattr(svc, "_clean_up_workflow_browser", AsyncMock(side_effect=cleanup_side_effect))
monkeypatch.setattr(
svc,
"_persist_workflow_browser_session_if_needed",
AsyncMock(side_effect=lambda **_kwargs: order.append("persist_helper")),
)
_patch_finalize(monkeypatch, svc, order, completed_run)
monkeypatch.setattr(svc, "persist_video_data", AsyncMock())
monkeypatch.setattr(svc, "execute_workflow_webhook", AsyncMock())
result = await _run_execute_workflow(svc)
assert result is completed_run
assert order[:2] == ["teardown_error", "finalize"]