From d0094d3568a75d11412759d54998fd34c7039ff4 Mon Sep 17 00:00:00 2001 From: Suchintan Date: Tue, 7 Jul 2026 17:45:11 -0400 Subject: [PATCH] feat(SKY-8207): consistent egress IP for Save & Reuse Session runs (#7169) Co-authored-by: Suchintan Singh Co-authored-by: AronPerez --- ...6_add_pin_saved_session_ip_to_workflows.py | 31 + docs/api-reference/openapi.json | 55 + .../WorkflowVisualComparisonDrawer.tsx | 1 + .../WorkflowCopilotChat.composerMode.test.tsx | 1 + .../WorkflowCopilotChat.queue.test.tsx | 1 + .../WorkflowCopilotChat.runGrounding.test.tsx | 1 + .../workflows/copilot/WorkflowCopilotChat.tsx | 2 + .../routes/workflows/debugger/Debugger.tsx | 1 + .../workflows/editor/WorkflowEditor.tsx | 1 + .../src/routes/workflows/editor/Workspace.tsx | 3 + .../editor/nestedConditionalSave.test.ts | 1 + .../editor/nodes/StartNode/StartNode.test.tsx | 1 + .../editor/nodes/StartNode/StartNode.tsx | 4 + .../StartNode/WorkflowSettingsEditor.test.tsx | 1 + .../StartNode/WorkflowSettingsEditor.tsx | 15 + .../workflows/editor/nodes/StartNode/types.ts | 1 + .../WorkflowSettingsBlockForm.tsx | 1 + .../editor/panels/WorkflowComparisonPanel.tsx | 1 + .../workflows/editor/sortable/perf.test.ts | 1 + .../editor/sortable/roundTrip.test.ts | 1 + .../src/routes/workflows/editor/utils.test.ts | 1 + .../workflows/editor/workflowEditorUtils.ts | 4 + .../editor/workflowVersionFromSaveData.ts | 1 + .../routes/workflows/types/workflowTypes.ts | 2 + .../workflows/types/workflowYamlTypes.ts | 1 + .../src/store/WorkflowHasChangesStore.ts | 1 + .../src/store/WorkflowSettingsStore.ts | 2 + skyvern/cli/mcp_tools/workflow.py | 2 +- skyvern/client/types/workflow.py | 1 + .../types/workflow_create_yaml_request.py | 1 + skyvern/forge/sdk/copilot/tools/frontier.py | 1 + .../sdk/copilot/tools/workflow_update.py | 1 + skyvern/forge/sdk/db/models.py | 1 + .../sdk/db/repositories/browser_sessions.py | 30 +- .../workflow_run_credential_selections.py | 13 + .../sdk/db/repositories/workflow_runs.py | 52 +- .../forge/sdk/db/repositories/workflows.py | 8 + skyvern/forge/sdk/db/utils.py | 1 + skyvern/forge/sdk/routes/workflow_copilot.py | 28 +- .../schemas/persistent_browser_sessions.py | 3 + skyvern/forge/sdk/workflow/models/workflow.py | 1 + skyvern/forge/sdk/workflow/service.py | 684 ++++++++-- skyvern/schemas/proxy_pinning.py | 7 + skyvern/schemas/workflows.py | 1 + .../default_persistent_sessions_manager.py | 2 + skyvern/webeye/persistent_sessions_manager.py | 1 + .../test_update_workflow_definition_fields.py | 12 +- .../sdk/db/test_workflow_runs_repository.py | 208 ++- .../test_copilot_code_block_persist_seam.py | 2 + tests/unit/test_credential_rotation.py | 192 +++ .../test_login_block_credential_proxy_pin.py | 16 +- tests/unit/test_saved_session_proxy_pin.py | 1111 +++++++++++++++++ .../unit/test_workflow_browser_profile_id.py | 49 + tests/unit/test_workflow_elapsed_timeout.py | 8 + ...orkflow_service_managed_browser_profile.py | 319 ++++- 55 files changed, 2745 insertions(+), 146 deletions(-) create mode 100644 alembic/versions/2026_07_07_2121-5b618ec6dce6_add_pin_saved_session_ip_to_workflows.py create mode 100644 tests/unit/test_saved_session_proxy_pin.py diff --git a/alembic/versions/2026_07_07_2121-5b618ec6dce6_add_pin_saved_session_ip_to_workflows.py b/alembic/versions/2026_07_07_2121-5b618ec6dce6_add_pin_saved_session_ip_to_workflows.py new file mode 100644 index 000000000..4ae5beb91 --- /dev/null +++ b/alembic/versions/2026_07_07_2121-5b618ec6dce6_add_pin_saved_session_ip_to_workflows.py @@ -0,0 +1,31 @@ +"""add pin_saved_session_ip to workflows + +Revision ID: 5b618ec6dce6 +Revises: e2d87251fee8 +Create Date: 2026-07-07T21:21:37.964227+00:00 + +""" + +from typing import Sequence, Union + +import sqlalchemy as sa + +from alembic import op + +# revision identifiers, used by Alembic. +revision: str = "5b618ec6dce6" +down_revision: Union[str, None] = "e2d87251fee8" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + op.execute("SET LOCAL lock_timeout = '5s'") + op.add_column( + "workflows", + sa.Column("pin_saved_session_ip", sa.Boolean(), nullable=False, server_default=sa.false()), + ) + + +def downgrade() -> None: + op.drop_column("workflows", "pin_saved_session_ip") diff --git a/docs/api-reference/openapi.json b/docs/api-reference/openapi.json index b8e5f02a9..43e821871 100644 --- a/docs/api-reference/openapi.json +++ b/docs/api-reference/openapi.json @@ -2683,6 +2683,24 @@ }, "description": "Case-insensitive substring search across: browser profile name and description. A profile is returned if either field matches." }, + { + "name": "managed", + "in": "query", + "required": false, + "schema": { + "anyOf": [ + { + "type": "boolean" + }, + { + "type": "null" + } + ], + "description": "Omit to return all profiles; false returns only user-created profiles; true returns only auto-managed profiles.", + "title": "Managed" + }, + "description": "Omit to return all profiles; false returns only user-created profiles; true returns only auto-managed profiles." + }, { "name": "x-api-key", "in": "header", @@ -8957,6 +8975,33 @@ ], "title": "Proxy Session Id" }, + "is_managed": { + "type": "boolean", + "title": "Is Managed", + "default": false + }, + "workflow_permanent_id": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Workflow Permanent Id" + }, + "browser_profile_key_digest": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "Browser Profile Key Digest" + }, "created_at": { "type": "string", "format": "date-time", @@ -23984,6 +24029,11 @@ "title": "Persist Browser Session", "default": false }, + "pin_saved_session_ip": { + "type": "boolean", + "title": "Pin Saved Session Ip", + "default": false + }, "browser_profile_id": { "anyOf": [ { @@ -24307,6 +24357,11 @@ "title": "Persist Browser Session", "default": false }, + "pin_saved_session_ip": { + "type": "boolean", + "title": "Pin Saved Session Ip", + "default": false + }, "browser_profile_id": { "anyOf": [ { diff --git a/skyvern-frontend/src/routes/workflows/components/WorkflowVisualComparisonDrawer.tsx b/skyvern-frontend/src/routes/workflows/components/WorkflowVisualComparisonDrawer.tsx index d58a131ea..4287c51d6 100644 --- a/skyvern-frontend/src/routes/workflows/components/WorkflowVisualComparisonDrawer.tsx +++ b/skyvern-frontend/src/routes/workflows/components/WorkflowVisualComparisonDrawer.tsx @@ -103,6 +103,7 @@ function getWorkflowElements(version: WorkflowVersion) { proxyLocation: version.proxy_location ?? ProxyLocation.Residential, webhookCallbackUrl: version.webhook_callback_url || "", persistBrowserSession: version.persist_browser_session, + pinSavedSessionIp: version.pin_saved_session_ip ?? false, browserProfileId: version.browser_profile_id ?? null, browserProfileKey: version.browser_profile_key ?? null, model: version.model, diff --git a/skyvern-frontend/src/routes/workflows/copilot/WorkflowCopilotChat.composerMode.test.tsx b/skyvern-frontend/src/routes/workflows/copilot/WorkflowCopilotChat.composerMode.test.tsx index 071f2aeb6..9c4ba8074 100644 --- a/skyvern-frontend/src/routes/workflows/copilot/WorkflowCopilotChat.composerMode.test.tsx +++ b/skyvern-frontend/src/routes/workflows/copilot/WorkflowCopilotChat.composerMode.test.tsx @@ -98,6 +98,7 @@ const saveData = { proxyLocation: null, webhookCallbackUrl: null, persistBrowserSession: false, + pinSavedSessionIp: false, browserProfileId: null, browserProfileKey: null, model: null, diff --git a/skyvern-frontend/src/routes/workflows/copilot/WorkflowCopilotChat.queue.test.tsx b/skyvern-frontend/src/routes/workflows/copilot/WorkflowCopilotChat.queue.test.tsx index ff6df0c3c..164f6b389 100644 --- a/skyvern-frontend/src/routes/workflows/copilot/WorkflowCopilotChat.queue.test.tsx +++ b/skyvern-frontend/src/routes/workflows/copilot/WorkflowCopilotChat.queue.test.tsx @@ -112,6 +112,7 @@ const saveData = { proxyLocation: null, webhookCallbackUrl: null, persistBrowserSession: false, + pinSavedSessionIp: false, browserProfileId: null, browserProfileKey: null, model: null, diff --git a/skyvern-frontend/src/routes/workflows/copilot/WorkflowCopilotChat.runGrounding.test.tsx b/skyvern-frontend/src/routes/workflows/copilot/WorkflowCopilotChat.runGrounding.test.tsx index 4c0eff527..a4a9b7125 100644 --- a/skyvern-frontend/src/routes/workflows/copilot/WorkflowCopilotChat.runGrounding.test.tsx +++ b/skyvern-frontend/src/routes/workflows/copilot/WorkflowCopilotChat.runGrounding.test.tsx @@ -101,6 +101,7 @@ const saveData = { proxyLocation: null, webhookCallbackUrl: null, persistBrowserSession: false, + pinSavedSessionIp: false, browserProfileId: null, browserProfileKey: null, model: null, diff --git a/skyvern-frontend/src/routes/workflows/copilot/WorkflowCopilotChat.tsx b/skyvern-frontend/src/routes/workflows/copilot/WorkflowCopilotChat.tsx index 97ebb36e4..64508d555 100644 --- a/skyvern-frontend/src/routes/workflows/copilot/WorkflowCopilotChat.tsx +++ b/skyvern-frontend/src/routes/workflows/copilot/WorkflowCopilotChat.tsx @@ -1376,6 +1376,7 @@ export function WorkflowCopilotChat({ proxy_location: saveData.settings.proxyLocation, webhook_callback_url: saveData.settings.webhookCallbackUrl, persist_browser_session: saveData.settings.persistBrowserSession, + pin_saved_session_ip: saveData.settings.pinSavedSessionIp, browser_profile_id: saveData.settings.browserProfileId, browser_profile_key: saveData.settings.browserProfileKey, model: saveData.settings.model, @@ -1415,6 +1416,7 @@ export function WorkflowCopilotChat({ proxy_location: saveData.settings.proxyLocation, webhook_callback_url: saveData.settings.webhookCallbackUrl, persist_browser_session: saveData.settings.persistBrowserSession, + pin_saved_session_ip: saveData.settings.pinSavedSessionIp, browser_profile_id: saveData.settings.browserProfileId, browser_profile_key: saveData.settings.browserProfileKey, model: saveData.settings.model, diff --git a/skyvern-frontend/src/routes/workflows/debugger/Debugger.tsx b/skyvern-frontend/src/routes/workflows/debugger/Debugger.tsx index af38ee440..b61b902bb 100644 --- a/skyvern-frontend/src/routes/workflows/debugger/Debugger.tsx +++ b/skyvern-frontend/src/routes/workflows/debugger/Debugger.tsx @@ -70,6 +70,7 @@ function Debugger() { const settings: WorkflowSettings = { persistBrowserSession: workflow.persist_browser_session, + pinSavedSessionIp: workflow.pin_saved_session_ip ?? false, browserProfileId: workflow.browser_profile_id ?? null, browserProfileKey: workflow.browser_profile_key ?? null, proxyLocation: workflow.proxy_location, diff --git a/skyvern-frontend/src/routes/workflows/editor/WorkflowEditor.tsx b/skyvern-frontend/src/routes/workflows/editor/WorkflowEditor.tsx index a1014a1c6..942c9dec5 100644 --- a/skyvern-frontend/src/routes/workflows/editor/WorkflowEditor.tsx +++ b/skyvern-frontend/src/routes/workflows/editor/WorkflowEditor.tsx @@ -99,6 +99,7 @@ function WorkflowEditor() { const settings: WorkflowSettings = { persistBrowserSession: workflow.persist_browser_session, + pinSavedSessionIp: workflow.pin_saved_session_ip ?? false, browserProfileId: workflow.browser_profile_id ?? null, browserProfileKey: workflow.browser_profile_key ?? null, proxyLocation: workflow.proxy_location, diff --git a/skyvern-frontend/src/routes/workflows/editor/Workspace.tsx b/skyvern-frontend/src/routes/workflows/editor/Workspace.tsx index 9b937eb42..186ef4973 100644 --- a/skyvern-frontend/src/routes/workflows/editor/Workspace.tsx +++ b/skyvern-frontend/src/routes/workflows/editor/Workspace.tsx @@ -1540,6 +1540,7 @@ function Workspace({ proxyLocation: workflowData.proxy_location ?? ProxyLocation.Residential, webhookCallbackUrl: workflowData.webhook_callback_url || "", persistBrowserSession: workflowData.persist_browser_session ?? false, + pinSavedSessionIp: workflowData.pin_saved_session_ip ?? false, browserProfileId: workflowData.browser_profile_id ?? null, browserProfileKey: workflowData.browser_profile_key ?? null, model: workflowData.model ?? null, @@ -1846,6 +1847,7 @@ function Workspace({ selectedVersion.proxy_location ?? ProxyLocation.Residential, webhookCallbackUrl: selectedVersion.webhook_callback_url || "", persistBrowserSession: selectedVersion.persist_browser_session, + pinSavedSessionIp: selectedVersion.pin_saved_session_ip ?? false, browserProfileId: selectedVersion.browser_profile_id ?? null, browserProfileKey: selectedVersion.browser_profile_key ?? null, model: selectedVersion.model, @@ -2867,6 +2869,7 @@ function Workspace({ extra_http_headers: extraHttpHeaders, cdp_connect_headers: cdpConnectHeaders, persist_browser_session: saveData.settings.persistBrowserSession, + pin_saved_session_ip: saveData.settings.pinSavedSessionIp, browser_profile_id: saveData.settings.browserProfileId, browser_profile_key: saveData.settings.browserProfileKey, model: saveData.settings.model, diff --git a/skyvern-frontend/src/routes/workflows/editor/nestedConditionalSave.test.ts b/skyvern-frontend/src/routes/workflows/editor/nestedConditionalSave.test.ts index 0ab9aca3e..5021d845f 100644 --- a/skyvern-frontend/src/routes/workflows/editor/nestedConditionalSave.test.ts +++ b/skyvern-frontend/src/routes/workflows/editor/nestedConditionalSave.test.ts @@ -20,6 +20,7 @@ const DEFAULT_SETTINGS: WorkflowSettings = { proxyLocation: ProxyLocation.Residential, webhookCallbackUrl: null, persistBrowserSession: false, + pinSavedSessionIp: false, browserProfileId: null, browserProfileKey: null, model: null, diff --git a/skyvern-frontend/src/routes/workflows/editor/nodes/StartNode/StartNode.test.tsx b/skyvern-frontend/src/routes/workflows/editor/nodes/StartNode/StartNode.test.tsx index f07738ec3..eb190b10e 100644 --- a/skyvern-frontend/src/routes/workflows/editor/nodes/StartNode/StartNode.test.tsx +++ b/skyvern-frontend/src/routes/workflows/editor/nodes/StartNode/StartNode.test.tsx @@ -46,6 +46,7 @@ const startNodeData: WorkflowStartNodeData = { webhookCallbackUrl: "", proxyLocation: null, persistBrowserSession: false, + pinSavedSessionIp: false, browserProfileId: null, browserProfileKey: null, model: null, diff --git a/skyvern-frontend/src/routes/workflows/editor/nodes/StartNode/StartNode.tsx b/skyvern-frontend/src/routes/workflows/editor/nodes/StartNode/StartNode.tsx index 9bb339f68..a3a5c1495 100644 --- a/skyvern-frontend/src/routes/workflows/editor/nodes/StartNode/StartNode.tsx +++ b/skyvern-frontend/src/routes/workflows/editor/nodes/StartNode/StartNode.tsx @@ -31,6 +31,7 @@ interface StartSettings { webhookCallbackUrl: string; proxyLocation: ProxyLocation; persistBrowserSession: boolean; + pinSavedSessionIp: boolean; browserProfileId: string | null; browserProfileKey: string | null; model: WorkflowModel | null; @@ -67,6 +68,9 @@ function StartNode({ id, data, parentId }: NodeProps) { persistBrowserSession: data.withWorkflowSettings ? data.persistBrowserSession : false, + pinSavedSessionIp: data.withWorkflowSettings + ? data.pinSavedSessionIp + : false, browserProfileId: data.withWorkflowSettings ? data.browserProfileId : null, diff --git a/skyvern-frontend/src/routes/workflows/editor/nodes/StartNode/WorkflowSettingsEditor.test.tsx b/skyvern-frontend/src/routes/workflows/editor/nodes/StartNode/WorkflowSettingsEditor.test.tsx index a882c0406..96bc2e2a3 100644 --- a/skyvern-frontend/src/routes/workflows/editor/nodes/StartNode/WorkflowSettingsEditor.test.tsx +++ b/skyvern-frontend/src/routes/workflows/editor/nodes/StartNode/WorkflowSettingsEditor.test.tsx @@ -13,6 +13,7 @@ const startNodeData: WorkflowStartNodeData = { webhookCallbackUrl: "", proxyLocation: null, persistBrowserSession: true, + pinSavedSessionIp: false, browserProfileId: null, browserProfileKey: null, model: null, diff --git a/skyvern-frontend/src/routes/workflows/editor/nodes/StartNode/WorkflowSettingsEditor.tsx b/skyvern-frontend/src/routes/workflows/editor/nodes/StartNode/WorkflowSettingsEditor.tsx index eb6bb47b6..a8da1e2f5 100644 --- a/skyvern-frontend/src/routes/workflows/editor/nodes/StartNode/WorkflowSettingsEditor.tsx +++ b/skyvern-frontend/src/routes/workflows/editor/nodes/StartNode/WorkflowSettingsEditor.tsx @@ -57,6 +57,9 @@ const SEQUENTIAL_KEY_TOOLTIP = const BROWSER_PROFILE_KEY_TOOLTIP = "Template for separating saved browser profiles. Use + to insert an agent input, or type a static key. Runs with the same rendered value reuse the same saved profile."; +const PIN_SAVED_SESSION_IP_TOOLTIP = + "Pin this workflow's saved sessions to a consistent proxy IP across runs, so restored logins are not invalidated by IP changes. Requires the Residential (ISP) proxy location. With a Browser Profile Key, each saved profile keeps its own IP."; + const WORKFLOW_RUN_DEFAULT_MAX_ELAPSED_TIME_MINUTES = 4 * 60; const WORKFLOW_RUN_MAX_ELAPSED_TIME_MINUTES = 8 * 60; @@ -324,6 +327,7 @@ function WorkflowSettingsEditorBody({ onCheckedChange={(value) => update({ persistBrowserSession: value, + pinSavedSessionIp: value ? data.pinSavedSessionIp : false, browserProfileKey: value ? data.browserProfileKey : null, }) } @@ -331,6 +335,17 @@ function WorkflowSettingsEditorBody({ {data.persistBrowserSession && (
+
+ + + + update({ pinSavedSessionIp: value }) + } + /> +
diff --git a/skyvern-frontend/src/routes/workflows/editor/nodes/StartNode/types.ts b/skyvern-frontend/src/routes/workflows/editor/nodes/StartNode/types.ts index 2510435d6..e8e6e4208 100644 --- a/skyvern-frontend/src/routes/workflows/editor/nodes/StartNode/types.ts +++ b/skyvern-frontend/src/routes/workflows/editor/nodes/StartNode/types.ts @@ -8,6 +8,7 @@ export type WorkflowStartNodeData = { webhookCallbackUrl: string; proxyLocation: ProxyLocation; persistBrowserSession: boolean; + pinSavedSessionIp: boolean; browserProfileId: string | null; browserProfileKey: string | null; model: WorkflowModel | null; diff --git a/skyvern-frontend/src/routes/workflows/editor/panels/BlockConfigForm/WorkflowSettingsBlockForm.tsx b/skyvern-frontend/src/routes/workflows/editor/panels/BlockConfigForm/WorkflowSettingsBlockForm.tsx index 8466e6b32..a67b73b27 100644 --- a/skyvern-frontend/src/routes/workflows/editor/panels/BlockConfigForm/WorkflowSettingsBlockForm.tsx +++ b/skyvern-frontend/src/routes/workflows/editor/panels/BlockConfigForm/WorkflowSettingsBlockForm.tsx @@ -39,6 +39,7 @@ function WorkflowSettingsBlockFormBody({ runSequentially: data.runSequentially, sequentialKey: data.sequentialKey, persistBrowserSession: data.persistBrowserSession, + pinSavedSessionIp: data.pinSavedSessionIp, extraHttpHeaders: data.extraHttpHeaders, maxScreenshotScrolls: data.maxScreenshotScrolls, maxElapsedTimeMinutes: data.maxElapsedTimeMinutes, diff --git a/skyvern-frontend/src/routes/workflows/editor/panels/WorkflowComparisonPanel.tsx b/skyvern-frontend/src/routes/workflows/editor/panels/WorkflowComparisonPanel.tsx index 4cf4f0cc1..6132eb2ea 100644 --- a/skyvern-frontend/src/routes/workflows/editor/panels/WorkflowComparisonPanel.tsx +++ b/skyvern-frontend/src/routes/workflows/editor/panels/WorkflowComparisonPanel.tsx @@ -142,6 +142,7 @@ function getWorkflowElements(version: WorkflowVersion) { proxyLocation: version.proxy_location ?? ProxyLocation.Residential, webhookCallbackUrl: version.webhook_callback_url || "", persistBrowserSession: version.persist_browser_session, + pinSavedSessionIp: version.pin_saved_session_ip ?? false, browserProfileId: version.browser_profile_id ?? null, browserProfileKey: version.browser_profile_key ?? null, model: version.model, diff --git a/skyvern-frontend/src/routes/workflows/editor/sortable/perf.test.ts b/skyvern-frontend/src/routes/workflows/editor/sortable/perf.test.ts index e40c95774..8d1373b56 100644 --- a/skyvern-frontend/src/routes/workflows/editor/sortable/perf.test.ts +++ b/skyvern-frontend/src/routes/workflows/editor/sortable/perf.test.ts @@ -49,6 +49,7 @@ const DEFAULT_SETTINGS: WorkflowSettings = { proxyLocation: ProxyLocation.Residential, webhookCallbackUrl: null, persistBrowserSession: false, + pinSavedSessionIp: false, browserProfileId: null, browserProfileKey: null, model: null, diff --git a/skyvern-frontend/src/routes/workflows/editor/sortable/roundTrip.test.ts b/skyvern-frontend/src/routes/workflows/editor/sortable/roundTrip.test.ts index b10f935ad..797547833 100644 --- a/skyvern-frontend/src/routes/workflows/editor/sortable/roundTrip.test.ts +++ b/skyvern-frontend/src/routes/workflows/editor/sortable/roundTrip.test.ts @@ -43,6 +43,7 @@ const DEFAULT_SETTINGS: WorkflowSettings = { proxyLocation: ProxyLocation.Residential, webhookCallbackUrl: null, persistBrowserSession: false, + pinSavedSessionIp: false, browserProfileId: null, browserProfileKey: null, model: null, diff --git a/skyvern-frontend/src/routes/workflows/editor/utils.test.ts b/skyvern-frontend/src/routes/workflows/editor/utils.test.ts index 88e1d15a9..3a50bcfbb 100644 --- a/skyvern-frontend/src/routes/workflows/editor/utils.test.ts +++ b/skyvern-frontend/src/routes/workflows/editor/utils.test.ts @@ -27,6 +27,7 @@ const baseWorkflow = { extra_http_headers: null, cdp_connect_headers: null, persist_browser_session: false, + pin_saved_session_ip: false, browser_profile_id: null, browser_profile_key: null, model: null, diff --git a/skyvern-frontend/src/routes/workflows/editor/workflowEditorUtils.ts b/skyvern-frontend/src/routes/workflows/editor/workflowEditorUtils.ts index 381b5a0b8..ebebd8f14 100644 --- a/skyvern-frontend/src/routes/workflows/editor/workflowEditorUtils.ts +++ b/skyvern-frontend/src/routes/workflows/editor/workflowEditorUtils.ts @@ -1994,6 +1994,7 @@ function getElements( startNode(startNodeId, { withWorkflowSettings: true, persistBrowserSession: settings.persistBrowserSession, + pinSavedSessionIp: settings.pinSavedSessionIp, browserProfileId: settings.browserProfileId, browserProfileKey: settings.browserProfileKey, proxyLocation: settings.proxyLocation ?? ProxyLocation.Residential, @@ -3416,6 +3417,7 @@ function getWorkflowBlocks( function getWorkflowSettings(nodes: Array): WorkflowSettings { const defaultSettings = { persistBrowserSession: false, + pinSavedSessionIp: false, browserProfileId: null, browserProfileKey: null, proxyLocation: ProxyLocation.Residential, @@ -3447,6 +3449,7 @@ function getWorkflowSettings(nodes: Array): WorkflowSettings { if (isWorkflowStartNodeData(data)) { return { persistBrowserSession: data.persistBrowserSession, + pinSavedSessionIp: data.pinSavedSessionIp, browserProfileId: data.browserProfileId, browserProfileKey: data.browserProfileKey, proxyLocation: data.proxyLocation, @@ -4542,6 +4545,7 @@ function convert(workflow: WorkflowApiResponse): WorkflowCreateYAMLRequest { proxy_location: workflow.proxy_location, webhook_callback_url: workflow.webhook_callback_url, persist_browser_session: workflow.persist_browser_session, + pin_saved_session_ip: workflow.pin_saved_session_ip, browser_profile_id: workflow.browser_profile_id ?? null, browser_profile_key: workflow.browser_profile_key ?? null, model: workflow.model, diff --git a/skyvern-frontend/src/routes/workflows/editor/workflowVersionFromSaveData.ts b/skyvern-frontend/src/routes/workflows/editor/workflowVersionFromSaveData.ts index 2c10ae4b9..fd03be692 100644 --- a/skyvern-frontend/src/routes/workflows/editor/workflowVersionFromSaveData.ts +++ b/skyvern-frontend/src/routes/workflows/editor/workflowVersionFromSaveData.ts @@ -48,6 +48,7 @@ export function workflowVersionFromSaveData( extra_http_headers: headers.extraHttpHeaders, cdp_connect_headers: headers.cdpConnectHeaders, persist_browser_session: settings.persistBrowserSession, + pin_saved_session_ip: settings.pinSavedSessionIp, browser_profile_id: settings.browserProfileId, browser_profile_key: settings.browserProfileKey, model: settings.model, diff --git a/skyvern-frontend/src/routes/workflows/types/workflowTypes.ts b/skyvern-frontend/src/routes/workflows/types/workflowTypes.ts index 91fb3f9cf..36670e6f0 100644 --- a/skyvern-frontend/src/routes/workflows/types/workflowTypes.ts +++ b/skyvern-frontend/src/routes/workflows/types/workflowTypes.ts @@ -666,6 +666,7 @@ export type WorkflowApiResponse = { extra_http_headers: Record | null; cdp_connect_headers: Record | null; persist_browser_session: boolean; + pin_saved_session_ip: boolean; browser_profile_id?: string | null; browser_profile_key?: string | null; model: WorkflowModel | null; @@ -696,6 +697,7 @@ export type WorkflowSettings = { proxyLocation: ProxyLocation | null; webhookCallbackUrl: string | null; persistBrowserSession: boolean; + pinSavedSessionIp: boolean; browserProfileId: string | null; browserProfileKey: string | null; model: WorkflowModel | null; diff --git a/skyvern-frontend/src/routes/workflows/types/workflowYamlTypes.ts b/skyvern-frontend/src/routes/workflows/types/workflowYamlTypes.ts index 23fb0f198..f64e8fb28 100644 --- a/skyvern-frontend/src/routes/workflows/types/workflowYamlTypes.ts +++ b/skyvern-frontend/src/routes/workflows/types/workflowYamlTypes.ts @@ -12,6 +12,7 @@ export type WorkflowCreateYAMLRequest = { proxy_location?: ProxyLocation | null; webhook_callback_url?: string | null; persist_browser_session?: boolean; + pin_saved_session_ip?: boolean; browser_profile_id?: string | null; browser_profile_key?: string | null; model?: WorkflowModel | null; diff --git a/skyvern-frontend/src/store/WorkflowHasChangesStore.ts b/skyvern-frontend/src/store/WorkflowHasChangesStore.ts index ed5067a87..cfff39b54 100644 --- a/skyvern-frontend/src/store/WorkflowHasChangesStore.ts +++ b/skyvern-frontend/src/store/WorkflowHasChangesStore.ts @@ -192,6 +192,7 @@ const useWorkflowSave = (opts?: WorkflowSaveOpts) => { proxy_location: saveData.settings.proxyLocation, webhook_callback_url: saveData.settings.webhookCallbackUrl, persist_browser_session: saveData.settings.persistBrowserSession, + pin_saved_session_ip: saveData.settings.pinSavedSessionIp, browser_profile_id: saveData.settings.browserProfileId, browser_profile_key: saveData.settings.browserProfileKey, model: saveData.settings.model, diff --git a/skyvern-frontend/src/store/WorkflowSettingsStore.ts b/skyvern-frontend/src/store/WorkflowSettingsStore.ts index 45fe44cd9..78f9b7019 100644 --- a/skyvern-frontend/src/store/WorkflowSettingsStore.ts +++ b/skyvern-frontend/src/store/WorkflowSettingsStore.ts @@ -9,6 +9,7 @@ export interface WorkflowSettingsState { webhookCallbackUrl: string; proxyLocation: ProxyLocation; persistBrowserSession: boolean; + pinSavedSessionIp: boolean; browserProfileKey: string | null; model: WorkflowModel | null; maxScreenshotScrollingTimes: number | null; @@ -32,6 +33,7 @@ const defaultState: Omit< webhookCallbackUrl: "", proxyLocation: ProxyLocation.Residential, persistBrowserSession: false, + pinSavedSessionIp: false, browserProfileKey: null, model: null, maxScreenshotScrollingTimes: null, diff --git a/skyvern/cli/mcp_tools/workflow.py b/skyvern/cli/mcp_tools/workflow.py index 64df10b23..516c9b978 100644 --- a/skyvern/cli/mcp_tools/workflow.py +++ b/skyvern/cli/mcp_tools/workflow.py @@ -1357,7 +1357,7 @@ async def skyvern_workflow_create( Pass run_with="code" to opt into cached script execution. Blocks share a browser session automatically. Leave optional toggles and overrides unset unless the user explicitly asks for them. This - applies to workflow-level fields (persist_browser_session, extra_http_headers, + applies to workflow-level fields (persist_browser_session, pin_saved_session_ip, extra_http_headers, totp_verification_url, totp_identifier, etc.) AND block-level overrides (max_retries, max_steps_per_run, totp_identifier, complete_criterion, error_code_mapping, continue_on_failure, engine, model, etc.). The schema defaults are intentional; silently diff --git a/skyvern/client/types/workflow.py b/skyvern/client/types/workflow.py index 648c0059d..b335adac1 100644 --- a/skyvern/client/types/workflow.py +++ b/skyvern/client/types/workflow.py @@ -27,6 +27,7 @@ class Workflow(UniversalBaseModel): totp_verification_url: typing.Optional[str] = None totp_identifier: typing.Optional[str] = None persist_browser_session: typing.Optional[bool] = None + pin_saved_session_ip: typing.Optional[bool] = None browser_profile_id: typing.Optional[str] = None browser_profile_key: typing.Optional[str] = None model: typing.Optional[typing.Dict[str, typing.Optional[typing.Any]]] = None diff --git a/skyvern/client/types/workflow_create_yaml_request.py b/skyvern/client/types/workflow_create_yaml_request.py index dd266183e..ef46b7f46 100644 --- a/skyvern/client/types/workflow_create_yaml_request.py +++ b/skyvern/client/types/workflow_create_yaml_request.py @@ -19,6 +19,7 @@ class WorkflowCreateYamlRequest(UniversalBaseModel): totp_verification_url: typing.Optional[str] = None totp_identifier: typing.Optional[str] = None persist_browser_session: typing.Optional[bool] = None + pin_saved_session_ip: typing.Optional[bool] = None browser_profile_id: typing.Optional[str] = None browser_profile_key: typing.Optional[str] = None model: typing.Optional[typing.Dict[str, typing.Optional[typing.Any]]] = None diff --git a/skyvern/forge/sdk/copilot/tools/frontier.py b/skyvern/forge/sdk/copilot/tools/frontier.py index 1411ca590..29cb06cc8 100644 --- a/skyvern/forge/sdk/copilot/tools/frontier.py +++ b/skyvern/forge/sdk/copilot/tools/frontier.py @@ -616,6 +616,7 @@ _CANONICAL_WORKFLOW_SETTING_FIELDS: tuple[str, ...] = ( "totp_verification_url", "totp_identifier", "persist_browser_session", + "pin_saved_session_ip", "browser_profile_id", "browser_profile_key", "model", diff --git a/skyvern/forge/sdk/copilot/tools/workflow_update.py b/skyvern/forge/sdk/copilot/tools/workflow_update.py index 0eef5d301..9652d8828 100644 --- a/skyvern/forge/sdk/copilot/tools/workflow_update.py +++ b/skyvern/forge/sdk/copilot/tools/workflow_update.py @@ -8647,6 +8647,7 @@ async def _update_workflow( totp_verification_url=workflow.totp_verification_url, totp_identifier=workflow.totp_identifier, persist_browser_session=workflow.persist_browser_session, + pin_saved_session_ip=workflow.pin_saved_session_ip, browser_profile_id=workflow.browser_profile_id, browser_profile_key=workflow.browser_profile_key, model=workflow.model, diff --git a/skyvern/forge/sdk/db/models.py b/skyvern/forge/sdk/db/models.py index c51bf9bec..de8e2317f 100644 --- a/skyvern/forge/sdk/db/models.py +++ b/skyvern/forge/sdk/db/models.py @@ -514,6 +514,7 @@ class WorkflowModel(SoftDeleteMixin, Base): totp_verification_url = Column(String) totp_identifier = Column(String) persist_browser_session = Column(Boolean, default=False, nullable=False) + pin_saved_session_ip = Column(Boolean, default=False, nullable=False, server_default=sqlalchemy.false()) browser_profile_id = Column(String, nullable=True) browser_profile_key = Column(String, nullable=True) model = Column(JSON, nullable=True) diff --git a/skyvern/forge/sdk/db/repositories/browser_sessions.py b/skyvern/forge/sdk/db/repositories/browser_sessions.py index a13c77fa2..bf3a2c661 100644 --- a/skyvern/forge/sdk/db/repositories/browser_sessions.py +++ b/skyvern/forge/sdk/db/repositories/browser_sessions.py @@ -2,6 +2,7 @@ from __future__ import annotations import uuid from datetime import datetime, timedelta +from typing import cast import structlog from sqlalchemy import case, desc, or_, select @@ -27,7 +28,7 @@ from skyvern.forge.sdk.schemas.persistent_browser_sessions import ( PersistentBrowserSession, PersistentBrowserType, ) -from skyvern.schemas.proxy_pinning import generate_proxy_session_id +from skyvern.schemas.proxy_pinning import generate_proxy_session_id, parse_proxy_location_input from skyvern.schemas.runs import ProxyLocation, ProxyLocationInput LOG = structlog.get_logger() @@ -387,17 +388,34 @@ class BrowserSessionsRepository(BaseRepository): browser_type: PersistentBrowserType | None = None, browser_profile_id: str | None = None, generate_browser_profile: bool = False, + inherit_profile_proxy: bool = False, ) -> PersistentBrowserSession: """Create a new persistent browser session.""" - proxy_location, proxy_session_id = normalize_proxy_pin_for_create( - proxy_location=proxy_location, - proxy_session_id=proxy_session_id, - ) extensions_str: list[str] | None = ( [extension.value for extension in extensions] if extensions is not None else None ) - serialized_proxy_location = serialize_proxy_location(proxy_location) async with self.Session() as session: + if inherit_profile_proxy and browser_profile_id and proxy_session_id is None: + query = ( + select(BrowserProfileModel) + .filter_by(browser_profile_id=browser_profile_id) + .filter_by(organization_id=organization_id) + .filter(BrowserProfileModel.deleted_at.is_(None)) + ) + browser_profile = (await session.scalars(query)).first() + if browser_profile and browser_profile.proxy_session_id: + proxy_session_id = browser_profile.proxy_session_id + # The ORM column stores the serialized string; deserialize before it flows + # into serialize_proxy_location, which rejects a bare str. + proxy_location = cast( + ProxyLocationInput, parse_proxy_location_input(browser_profile.proxy_location) + ) + + proxy_location, proxy_session_id = normalize_proxy_pin_for_create( + proxy_location=proxy_location, + proxy_session_id=proxy_session_id, + ) + serialized_proxy_location = serialize_proxy_location(proxy_location) browser_session = PersistentBrowserSessionModel( organization_id=organization_id, runnable_type=runnable_type, diff --git a/skyvern/forge/sdk/db/repositories/workflow_run_credential_selections.py b/skyvern/forge/sdk/db/repositories/workflow_run_credential_selections.py index d0b0cb608..72e1d9cc5 100644 --- a/skyvern/forge/sdk/db/repositories/workflow_run_credential_selections.py +++ b/skyvern/forge/sdk/db/repositories/workflow_run_credential_selections.py @@ -23,6 +23,19 @@ class WorkflowRunCredentialSelectionsRepository(BaseRepository): ).first() return selection.credential_id if selection else None + @db_operation("get_workflow_run_credential_selections_for_run") + async def get_selections_for_run(self, workflow_run_id: str) -> dict[str, str]: + async with self.Session() as session: + rows = ( + await session.execute( + select( + WorkflowRunCredentialSelectionModel.parameter_key, + WorkflowRunCredentialSelectionModel.credential_id, + ).where(WorkflowRunCredentialSelectionModel.workflow_run_id == workflow_run_id) + ) + ).all() + return {parameter_key: credential_id for parameter_key, credential_id in rows} + async def _get_selection(self, session: AsyncSession, workflow_run_id: str, parameter_key: str) -> str | None: selection = ( await session.scalars( diff --git a/skyvern/forge/sdk/db/repositories/workflow_runs.py b/skyvern/forge/sdk/db/repositories/workflow_runs.py index af7cff558..74ef678f6 100644 --- a/skyvern/forge/sdk/db/repositories/workflow_runs.py +++ b/skyvern/forge/sdk/db/repositories/workflow_runs.py @@ -35,6 +35,7 @@ if TYPE_CHECKING: from skyvern.forge.sdk.db._sentinels import _UNSET from skyvern.forge.sdk.db.models import ( + PersistentBrowserSessionModel, TaskModel, TaskRunModel, WorkflowModel, @@ -54,6 +55,7 @@ from skyvern.forge.sdk.db.utils import ( truncate_oversized_jsonb_value, ) from skyvern.forge.sdk.log_artifacts import save_workflow_run_logs +from skyvern.forge.sdk.schemas.persistent_browser_sessions import FORCED_WORKFLOW_SESSION_RUNNABLE_TYPE from skyvern.forge.sdk.schemas.tasks import Task from skyvern.forge.sdk.workflow.models.parameter import WorkflowParameter from skyvern.forge.sdk.workflow.models.workflow import ( @@ -726,10 +728,12 @@ class WorkflowRunsRepository(BaseRepository): workflow_permanent_id: str, organization_id: str | None = None, sequential_key: str | None = None, + include_browser_session_rows: bool = False, ) -> WorkflowRun | None: async with self.Session() as session: query = select(WorkflowRunModel).filter_by(workflow_permanent_id=workflow_permanent_id) - query = query.filter(WorkflowRunModel.browser_session_id.is_(None)) + if not include_browser_session_rows: + query = query.filter(WorkflowRunModel.browser_session_id.is_(None)) if organization_id: query = query.filter_by(organization_id=organization_id) query = query.filter_by(status=WorkflowRunStatus.queued) @@ -761,10 +765,12 @@ class WorkflowRunsRepository(BaseRepository): workflow_permanent_id: str, organization_id: str | None = None, sequential_key: str | None = None, + include_browser_session_rows: bool = False, ) -> WorkflowRun | None: async with self.Session() as session: query = select(WorkflowRunModel).filter_by(workflow_permanent_id=workflow_permanent_id) - query = query.filter(WorkflowRunModel.browser_session_id.is_(None)) + if not include_browser_session_rows: + query = query.filter(WorkflowRunModel.browser_session_id.is_(None)) if organization_id: query = query.filter_by(organization_id=organization_id) query = query.filter_by(status=WorkflowRunStatus.running) @@ -792,11 +798,35 @@ class WorkflowRunsRepository(BaseRepository): if run is None: return None - # Lane resolution mirrors enqueue priority: browser_session_id > browser_address - # > sequential_key > whole workflow. A debug-session run carries a browser_session_id - # but is excluded from the session lane (as at enqueue), so it serializes on its key. - query = select(WorkflowRunModel).filter_by(organization_id=run.organization_id) + self_forced = False if run.browser_session_id and not run.debug_session_id: + try: + persistent_browser_session = ( + await session.scalars( + select(PersistentBrowserSessionModel).filter_by( + persistent_browser_session_id=run.browser_session_id, + organization_id=run.organization_id, + ) + ) + ).first() + self_forced = ( + persistent_browser_session is not None + and persistent_browser_session.runnable_type == FORCED_WORKFLOW_SESSION_RUNNABLE_TYPE + ) + except Exception: + LOG.warning( + "Failed to fetch persistent browser session for runtime lane selection", + workflow_run_id=workflow_run_id, + browser_session_id=run.browser_session_id, + organization_id=run.organization_id, + exc_info=True, + ) + + # Lane resolution mirrors enqueue priority: browser_session_id > browser_address + # > sequential_key > whole workflow. Debug and forced-session runs carry a + # browser_session_id but are excluded from the session lane as they are at enqueue. + query = select(WorkflowRunModel).filter_by(organization_id=run.organization_id) + if run.browser_session_id and not run.debug_session_id and not self_forced: query = query.filter_by(browser_session_id=run.browser_session_id) elif run.browser_address: query = query.filter_by(browser_address=run.browser_address) @@ -804,11 +834,13 @@ class WorkflowRunsRepository(BaseRepository): query = query.filter_by( workflow_permanent_id=run.workflow_permanent_id, sequential_key=run.sequential_key, - ).filter(WorkflowRunModel.browser_session_id.is_(None)) - else: - query = query.filter_by(workflow_permanent_id=run.workflow_permanent_id).filter( - WorkflowRunModel.browser_session_id.is_(None) ) + if not self_forced: + query = query.filter(WorkflowRunModel.browser_session_id.is_(None)) + else: + query = query.filter_by(workflow_permanent_id=run.workflow_permanent_id) + if not self_forced: + query = query.filter(WorkflowRunModel.browser_session_id.is_(None)) # Sequential runs are stamped queued_at before Temporal submission; the fallback # only guards hand-created rows (e.g. tests) from comparing against None. diff --git a/skyvern/forge/sdk/db/repositories/workflows.py b/skyvern/forge/sdk/db/repositories/workflows.py index a4e1aec3c..6b8317d8a 100644 --- a/skyvern/forge/sdk/db/repositories/workflows.py +++ b/skyvern/forge/sdk/db/repositories/workflows.py @@ -102,6 +102,7 @@ class WorkflowsRepository(BaseRepository): totp_verification_url: str | None = None, totp_identifier: str | None = None, persist_browser_session: bool = False, + pin_saved_session_ip: bool = False, browser_profile_id: str | None = None, browser_profile_key: str | None = None, model: dict[str, Any] | None = None, @@ -137,6 +138,7 @@ class WorkflowsRepository(BaseRepository): extra_http_headers=extra_http_headers, cdp_connect_headers=cdp_connect_headers, persist_browser_session=persist_browser_session, + pin_saved_session_ip=pin_saved_session_ip, browser_profile_id=browser_profile_id, browser_profile_key=browser_profile_key, model=model, @@ -670,6 +672,7 @@ class WorkflowsRepository(BaseRepository): totp_verification_url: str | None | object = _UNSET, totp_identifier: str | None | object = _UNSET, persist_browser_session: bool | None = None, + pin_saved_session_ip: bool | None = None, browser_profile_id: str | None | object = _UNSET, browser_profile_key: str | None | object = _UNSET, model: dict[str, Any] | None | object = _UNSET, @@ -720,6 +723,8 @@ class WorkflowsRepository(BaseRepository): workflow.totp_identifier = cast(str | None, totp_identifier) if persist_browser_session is not None: workflow.persist_browser_session = persist_browser_session + if pin_saved_session_ip is not None: + workflow.pin_saved_session_ip = pin_saved_session_ip if browser_profile_id is not _UNSET: workflow.browser_profile_id = cast(str | None, browser_profile_id) if browser_profile_key is not _UNSET: @@ -986,6 +991,7 @@ class WorkflowsRepository(BaseRepository): totp_verification_url: str | None | object = _UNSET, totp_identifier: str | None | object = _UNSET, persist_browser_session: bool | None = None, + pin_saved_session_ip: bool | None = None, browser_profile_id: str | None | object = _UNSET, browser_profile_key: str | None | object = _UNSET, model: dict[str, Any] | None | object = _UNSET, @@ -1072,6 +1078,8 @@ class WorkflowsRepository(BaseRepository): workflow.totp_identifier = cast(str | None, totp_identifier) if persist_browser_session is not None: workflow.persist_browser_session = persist_browser_session + if pin_saved_session_ip is not None: + workflow.pin_saved_session_ip = pin_saved_session_ip if browser_profile_id is not _UNSET: workflow.browser_profile_id = cast(str | None, browser_profile_id) if browser_profile_key is not _UNSET: diff --git a/skyvern/forge/sdk/db/utils.py b/skyvern/forge/sdk/db/utils.py index c389b7ef2..443a83dfc 100644 --- a/skyvern/forge/sdk/db/utils.py +++ b/skyvern/forge/sdk/db/utils.py @@ -559,6 +559,7 @@ def convert_to_workflow( totp_verification_url=workflow_model.totp_verification_url, totp_identifier=workflow_model.totp_identifier, persist_browser_session=workflow_model.persist_browser_session, + pin_saved_session_ip=workflow_model.pin_saved_session_ip, browser_profile_id=workflow_model.browser_profile_id, browser_profile_key=workflow_model.browser_profile_key, model=workflow_model.model, diff --git a/skyvern/forge/sdk/routes/workflow_copilot.py b/skyvern/forge/sdk/routes/workflow_copilot.py index ca1873190..960a333ff 100644 --- a/skyvern/forge/sdk/routes/workflow_copilot.py +++ b/skyvern/forge/sdk/routes/workflow_copilot.py @@ -932,6 +932,7 @@ async def _commit_staged_workflow( totp_verification_url=staged_workflow.totp_verification_url, totp_identifier=staged_workflow.totp_identifier, persist_browser_session=staged_workflow.persist_browser_session, + pin_saved_session_ip=staged_workflow.pin_saved_session_ip, browser_profile_id=staged_workflow.browser_profile_id, browser_profile_key=staged_workflow.browser_profile_key, model=staged_workflow.model, @@ -970,6 +971,7 @@ async def _restore_workflow_definition(original_workflow: Workflow | None, organ totp_verification_url=original_workflow.totp_verification_url, totp_identifier=original_workflow.totp_identifier, persist_browser_session=original_workflow.persist_browser_session, + pin_saved_session_ip=original_workflow.pin_saved_session_ip, browser_profile_id=original_workflow.browser_profile_id, browser_profile_key=original_workflow.browser_profile_key, model=original_workflow.model, @@ -1567,7 +1569,7 @@ def _normalize_copilot_yaml(workflow_yaml: str) -> WorkflowCreateYAMLRequest: return workflow_yaml_request -def _yaml_enable_self_healing(workflow_yaml: str | None) -> bool | None: +def _yaml_bool_setting(workflow_yaml: str | None, setting_name: str) -> bool | None: if not workflow_yaml: return None try: @@ -1576,10 +1578,18 @@ def _yaml_enable_self_healing(workflow_yaml: str | None) -> bool | None: return None if not isinstance(parsed, dict): return None - value = parsed.get("enable_self_healing") + value = parsed.get(setting_name) return value if isinstance(value, bool) else None +def _yaml_enable_self_healing(workflow_yaml: str | None) -> bool | None: + return _yaml_bool_setting(workflow_yaml, "enable_self_healing") + + +def _yaml_pin_saved_session_ip(workflow_yaml: str | None) -> bool | None: + return _yaml_bool_setting(workflow_yaml, "pin_saved_session_ip") + + async def _process_workflow_yaml( workflow_id: str, workflow_permanent_id: str, @@ -1607,7 +1617,13 @@ async def _process_workflow_yaml( # unsaved editor toggle survives an unrelated copilot edit. A persisted-lookup # failure propagates: failing the save is safer than writing an implicit disable. enable_self_healing = _yaml_enable_self_healing(settings_fallback_yaml) - if enable_self_healing is None: + + pin_saved_session_ip = _yaml_pin_saved_session_ip(workflow_yaml) + if pin_saved_session_ip is None: + pin_saved_session_ip = _yaml_pin_saved_session_ip(settings_fallback_yaml) + + current_workflow: Workflow | None = None + if enable_self_healing is None or pin_saved_session_ip is None: try: current_workflow = await app.WORKFLOW_SERVICE.get_workflow_by_permanent_id( workflow_permanent_id=workflow_permanent_id, @@ -1615,7 +1631,10 @@ async def _process_workflow_yaml( ) except WorkflowNotFound: current_workflow = None - enable_self_healing = bool(current_workflow and current_workflow.enable_self_healing) + if enable_self_healing is None: + enable_self_healing = bool(current_workflow and getattr(current_workflow, "enable_self_healing", False)) + if pin_saved_session_ip is None: + pin_saved_session_ip = bool(current_workflow and getattr(current_workflow, "pin_saved_session_ip", False)) now = datetime.now(timezone.utc) return Workflow( @@ -1632,6 +1651,7 @@ async def _process_workflow_yaml( totp_verification_url=workflow_yaml_request.totp_verification_url, totp_identifier=workflow_yaml_request.totp_identifier, persist_browser_session=workflow_yaml_request.persist_browser_session or False, + pin_saved_session_ip=pin_saved_session_ip, browser_profile_id=workflow_yaml_request.browser_profile_id, browser_profile_key=workflow_yaml_request.browser_profile_key, model=workflow_yaml_request.model, diff --git a/skyvern/forge/sdk/schemas/persistent_browser_sessions.py b/skyvern/forge/sdk/schemas/persistent_browser_sessions.py index b297d7025..58054db0a 100644 --- a/skyvern/forge/sdk/schemas/persistent_browser_sessions.py +++ b/skyvern/forge/sdk/schemas/persistent_browser_sessions.py @@ -60,6 +60,9 @@ class Extensions(StrEnum): CaptchaSolver = "captcha-solver" +FORCED_WORKFLOW_SESSION_RUNNABLE_TYPE = "forced_workflow_run" + + class PersistentBrowserSession(BaseModel): model_config = ConfigDict(from_attributes=True) diff --git a/skyvern/forge/sdk/workflow/models/workflow.py b/skyvern/forge/sdk/workflow/models/workflow.py index e14a1de57..ee5d91399 100644 --- a/skyvern/forge/sdk/workflow/models/workflow.py +++ b/skyvern/forge/sdk/workflow/models/workflow.py @@ -137,6 +137,7 @@ class Workflow(BaseModel): totp_verification_url: str | None = None totp_identifier: str | None = None persist_browser_session: bool = False + pin_saved_session_ip: bool = False browser_profile_id: str | None = None browser_profile_key: str | None = None model: dict[str, Any] | None = None diff --git a/skyvern/forge/sdk/workflow/service.py b/skyvern/forge/sdk/workflow/service.py index c19f99d66..901a0517b 100644 --- a/skyvern/forge/sdk/workflow/service.py +++ b/skyvern/forge/sdk/workflow/service.py @@ -66,9 +66,13 @@ from skyvern.forge.sdk.db.id import generate_output_parameter_id, generate_workf from skyvern.forge.sdk.enterprise_features import collect_enterprise_gated_run_features from skyvern.forge.sdk.experimentation.enrich_tree import resolve_enrich_tree_for_context from skyvern.forge.sdk.models import Step, StepStatus +from skyvern.forge.sdk.schemas.browser_profiles import BrowserProfile from skyvern.forge.sdk.schemas.files import FileInfo from skyvern.forge.sdk.schemas.organizations import Organization -from skyvern.forge.sdk.schemas.persistent_browser_sessions import PersistentBrowserSession +from skyvern.forge.sdk.schemas.persistent_browser_sessions import ( + FORCED_WORKFLOW_SESSION_RUNNABLE_TYPE, + PersistentBrowserSession, +) from skyvern.forge.sdk.schemas.tasks import Task, TaskStatus from skyvern.forge.sdk.schemas.workflow_runs import WorkflowRunBlock, WorkflowRunTimeline, WorkflowRunTimelineType from skyvern.forge.sdk.trace import traced @@ -138,6 +142,11 @@ from skyvern.forge.sdk.workflow.status_mapping import ( TASK_STATUS_MAP, ) from skyvern.forge.sdk.workflow.workflow_definition_converter import convert_workflow_definition +from skyvern.schemas.proxy_pinning import ( + derive_proxy_session_id, + redact_proxy_session_id, + should_generate_proxy_session_id, +) from skyvern.schemas.run_enums import RunEngine from skyvern.schemas.runs import ( ProxyLocation, @@ -188,6 +197,9 @@ RECORDING_WINDOW_END_BUFFER = timedelta(minutes=15) # Skip post-run work when only a sub-millisecond budget remains; asyncio.timeout would fire on the first await. POST_RUN_TIMEOUT_EXHAUSTED_THRESHOLD_SECONDS = 0.001 +# Bound pre-finalization write-back so a stuck profile upload cannot leave the run in a non-terminal state. +BROWSER_SESSION_WRITE_BACK_TIMEOUT = SAVE_DOWNLOADED_FILES_TIMEOUT + # Short lifespan for auto-provisioned CodeBlock sessions; the renewal loop extends it while the # run is active, so this only bounds how long a leaked session lingers if cleanup never runs. CODE_BLOCK_SESSION_TIMEOUT_MINUTES = 20 @@ -431,6 +443,16 @@ class CachedScriptBlocks: blocks_to_clear: list[ScriptBlock] +@dataclass +class WorkflowBrowserCleanupResult: + browser_state: BrowserState | None + tasks: list[Task] + all_workflow_task_ids: list[str] + child_workflow_run_ids: list[str] + close_browser_on_completion: bool + browser_session_write_back_attempted: bool = False + + def _get_workflow_definition_core_data(workflow_definition: WorkflowDefinition) -> dict[str, Any]: """ This function dumps the workflow definition and removes the irrelevant data to the definition, like created_at and modified_at fields inside: @@ -1020,6 +1042,50 @@ class WorkflowService: await self._validate_credential_ids(credential_ids_to_validate, organization) + @staticmethod + def _get_rotating_credential_parameters(workflow: Workflow) -> list[CredentialParameter]: + workflow_definition = getattr(workflow, "workflow_definition", None) + if workflow_definition is None: + return [] + return [ + parameter + for parameter in workflow_definition.parameters + if isinstance(parameter, CredentialParameter) and bool(parameter.credential_ids) + ] + + async def _select_rotating_credential_parameters_for_render( + self, + *, + workflow: Workflow, + workflow_run: WorkflowRun, + organization_id: str, + ) -> dict[str, str]: + try: + selections: dict[str, str] = {} + for parameter in self._get_rotating_credential_parameters(workflow): + credential_ids = parameter.credential_ids + if not credential_ids: + continue + selections[parameter.key] = await select_credential_for_run( + workflow_run_id=workflow_run.workflow_run_id, + organization_id=organization_id, + workflow_permanent_id=workflow.workflow_permanent_id, + parameter_key=parameter.key, + credential_ids=credential_ids, + selection_strategy=parameter.selection_strategy, + ) + return selections + except Exception: + LOG.warning( + "Failed to select rotating credentials for workflow render parameters", + workflow_run_id=workflow_run.workflow_run_id, + workflow_permanent_id=workflow.workflow_permanent_id, + exc_info=True, + ) + if workflow.browser_profile_key: + raise + return {} + async def validate_schedule_parameters( self, workflow: Workflow, @@ -1336,10 +1402,21 @@ class WorkflowService: organization, ) + parameter_values = {param.key: value for param, value in workflow_parameter_values} + rotating_credential_selections = await self._select_rotating_credential_parameters_for_render( + workflow=workflow, + workflow_run=workflow_run, + organization_id=organization.organization_id, + ) + parameter_values.update(rotating_credential_selections) workflow_run = await self._prepare_persisted_workflow_browser_profile( workflow=workflow, workflow_run=workflow_run, - parameter_values={param.key: value for param, value in workflow_parameter_values}, + parameter_values=parameter_values, + # Keyless workflows keep best-effort rotation selection; keyed workflows re-raise above. + allow_missing_browser_profile_key=( + bool(self._get_rotating_credential_parameters(workflow)) and not rotating_credential_selections + ), ) if workflow_parameter_values: @@ -1407,6 +1484,7 @@ class WorkflowService: workflow: Workflow, workflow_run: WorkflowRun, parameter_values: dict[str, Any], + allow_missing_browser_profile_key: bool = False, ) -> WorkflowRun: if not workflow.persist_browser_session: return workflow_run @@ -1414,11 +1492,22 @@ class WorkflowService: if workflow_run.browser_profile_id: return workflow_run - rendered_key = await self._render_workflow_browser_profile_key( - workflow=workflow, - workflow_run=workflow_run, - parameter_values=parameter_values, - ) + try: + rendered_key = await self._render_workflow_browser_profile_key( + workflow=workflow, + workflow_run=workflow_run, + parameter_values=parameter_values, + ) + except MissingValueForParameter: + if not allow_missing_browser_profile_key: + raise + LOG.warning( + "Falling back to keyless managed browser profile after missing browser profile key", + workflow_run_id=workflow_run.workflow_run_id, + workflow_permanent_id=workflow.workflow_permanent_id, + browser_profile_key=workflow.browser_profile_key, + ) + rendered_key = None digest = build_browser_profile_key_digest(rendered_key) profile, created = await app.DATABASE.browser_sessions.get_or_create_managed_browser_profile( organization_id=workflow_run.organization_id, @@ -1430,7 +1519,8 @@ class WorkflowService: try: await self._seed_managed_browser_profile_from_legacy_session( workflow=workflow, - workflow_run=workflow_run, + organization_id=workflow_run.organization_id, + workflow_run_id=workflow_run.workflow_run_id, rendered_key=rendered_key, browser_profile_id=profile.browser_profile_id, ) @@ -1448,6 +1538,15 @@ class WorkflowService: organization_id=workflow_run.organization_id, ) return workflow_run + + await self._reconcile_managed_browser_profile_proxy_pin( + workflow=workflow, + profile=profile, + browser_profile_key_digest=digest, + organization_id=workflow_run.organization_id, + proxy_location=workflow_run.proxy_location, + workflow_run_id=workflow_run.workflow_run_id, + ) updated_workflow_run = await app.DATABASE.workflow_runs.update_workflow_run( workflow_run_id=workflow_run.workflow_run_id, browser_profile_id=profile.browser_profile_id, @@ -1461,29 +1560,175 @@ class WorkflowService: ) return updated_workflow_run + async def _reconcile_managed_browser_profile_proxy_pin( + self, + *, + workflow: Workflow, + profile: BrowserProfile, + browser_profile_key_digest: str, + organization_id: str, + proxy_location: ProxyLocationInput, + workflow_run_id: str | None = None, + ) -> None: + should_pin = workflow.pin_saved_session_ip and should_generate_proxy_session_id(proxy_location) + try: + existing_pin = profile.proxy_session_id + if should_pin: + if browser_profile_key_digest: + proxy_session_id = derive_proxy_session_id( + organization_id, + workflow.workflow_permanent_id, + browser_profile_key_digest, + ) + else: + proxy_session_id = derive_proxy_session_id( + organization_id, + workflow.workflow_permanent_id, + ) + # Re-derive rather than trust existing values so a drifted or corrupted + # pin or location self-heals to the deterministic per-segment state. + if existing_pin == proxy_session_id and should_generate_proxy_session_id(profile.proxy_location): + return + await app.DATABASE.browser_sessions.update_browser_profile( + profile_id=profile.browser_profile_id, + organization_id=organization_id, + proxy_location=ProxyLocation.RESIDENTIAL_ISP, + proxy_session_id=proxy_session_id, + ) + LOG.info( + "Pinned managed browser profile proxy session", + workflow_run_id=workflow_run_id, + workflow_permanent_id=workflow.workflow_permanent_id, + browser_profile_id=profile.browser_profile_id, + proxy_session_id=redact_proxy_session_id(proxy_session_id), + ) + return + + if profile.is_managed and existing_pin: + await app.DATABASE.browser_sessions.update_browser_profile( + profile_id=profile.browser_profile_id, + organization_id=organization_id, + proxy_location=None, + proxy_session_id=None, + ) + LOG.info( + "Cleared managed browser profile proxy session pin", + workflow_run_id=workflow_run_id, + workflow_permanent_id=workflow.workflow_permanent_id, + browser_profile_id=profile.browser_profile_id, + proxy_session_id=redact_proxy_session_id(existing_pin), + ) + except Exception: + LOG.warning( + "Failed to reconcile managed browser profile proxy pin", + workflow_run_id=workflow_run_id, + workflow_permanent_id=workflow.workflow_permanent_id, + browser_profile_id=profile.browser_profile_id, + exc_info=True, + ) + if should_pin: + raise + + async def _resolve_managed_browser_profile_for_run_request( + self, + *, + workflow: Workflow, + organization_id: str, + workflow_request: WorkflowRequestBody, + effective_proxy_location: ProxyLocationInput, + extra_parameter_values: dict[str, str] | None = None, + ) -> str | None: + try: + if not workflow.persist_browser_session or workflow_request.browser_profile_id: + return None + + rendered_key = None + if workflow.browser_profile_key: + parameter_values: dict[str, Any] = {} + for parameter in workflow.workflow_definition.parameters: + if isinstance(parameter, WorkflowParameter) and parameter.default_value is not None: + parameter_values[parameter.key] = parameter.default_value + if workflow_request.data: + parameter_values.update( + {key: value for key, value in workflow_request.data.items() if value is not None} + ) + if extra_parameter_values: + parameter_values.update(extra_parameter_values) + rendered_key = render_browser_profile_key(workflow.browser_profile_key, parameter_values) + if not rendered_key: + return None + + digest = build_browser_profile_key_digest(rendered_key) + profile, created = await app.DATABASE.browser_sessions.get_or_create_managed_browser_profile( + organization_id=organization_id, + workflow_permanent_id=workflow.workflow_permanent_id, + browser_profile_key_digest=digest, + name=_build_managed_browser_profile_name(workflow.title, rendered_key), + ) + if created: + try: + await self._seed_managed_browser_profile_from_legacy_session( + workflow=workflow, + organization_id=organization_id, + rendered_key=rendered_key, + browser_profile_id=profile.browser_profile_id, + ) + except Exception: + # Roll back the empty row so the next run re-attempts creation + seed; + # this run degrades to the legacy session path instead of failing setup. + LOG.warning( + "Failed to seed managed browser profile from legacy session; rolling back and using legacy path", + workflow_run_id=None, + browser_profile_id=profile.browser_profile_id, + exc_info=True, + ) + await app.DATABASE.browser_sessions.hard_delete_browser_profile( + profile_id=profile.browser_profile_id, + organization_id=organization_id, + ) + return None + await self._reconcile_managed_browser_profile_proxy_pin( + workflow=workflow, + profile=profile, + browser_profile_key_digest=digest, + organization_id=organization_id, + proxy_location=effective_proxy_location, + workflow_run_id=None, + ) + return profile.browser_profile_id + except Exception: + LOG.warning( + "Failed to resolve managed browser profile for workflow run request", + organization_id=organization_id, + workflow_permanent_id=workflow.workflow_permanent_id, + exc_info=True, + ) + return None + async def _seed_managed_browser_profile_from_legacy_session( self, *, workflow: Workflow, - workflow_run: WorkflowRun, + organization_id: str, rendered_key: str | None, browser_profile_id: str, + workflow_run_id: str | None = None, ) -> None: # Lazy migration: carry a workflow's accumulated Save & Reuse login state onto its new # managed profile the first time we create it, so the next run starts logged in. The OSS # launch path has no legacy-archive fallback, so this copy must happen here, not at launch. storage_key = build_workflow_browser_session_storage_key(workflow.workflow_permanent_id, rendered_key) - legacy_session_dir = await app.STORAGE.retrieve_browser_session(workflow_run.organization_id, storage_key) + legacy_session_dir = await app.STORAGE.retrieve_browser_session(organization_id, storage_key) if not legacy_session_dir: return await app.STORAGE.store_browser_profile( - workflow_run.organization_id, + organization_id, profile_id=browser_profile_id, directory=legacy_session_dir, ) LOG.info( "Seeded managed browser profile from legacy session archive", - workflow_run_id=workflow_run.workflow_run_id, + workflow_run_id=workflow_run_id, browser_profile_id=browser_profile_id, browser_session_storage_key=storage_key, ) @@ -1635,6 +1880,7 @@ class WorkflowService: timeout_minutes=timeout_seconds // 60, browser_profile_id=browser_profile_id, proxy_location=proxy_location, + inherit_profile_proxy=True, ) return browser_session @@ -1692,6 +1938,7 @@ class WorkflowService: timeout_minutes=CODE_BLOCK_SESSION_TIMEOUT_MINUTES, browser_profile_id=browser_profile_id, proxy_location=proxy_location, + inherit_profile_proxy=True, ) except Exception: LOG.warning( @@ -2332,7 +2579,46 @@ class WorkflowService: # in place of the real terminal reason. When pre_finally_status is # still ``None`` (cancellation landed before block execution # completed), there's no captured intent to restore and we skip. + browser_cleanup_result: WorkflowBrowserCleanupResult | None = None if pre_finally_status is not None: + if workflow.persist_browser_session and pre_finally_status not in ( + WorkflowRunStatus.canceled, + WorkflowRunStatus.failed, + WorkflowRunStatus.terminated, + WorkflowRunStatus.timed_out, + ): + try: + # Sequential dependency gates clear on terminal status, so persisted profile state must land first. + browser_cleanup_result = await self._clean_up_workflow_browser( + workflow_run=workflow_run, + close_browser_on_completion=close_browser_on_completion, + browser_session_id=browser_session_id, + ) + except BaseException: + LOG.warning( + "Pre-finalization browser cleanup failed during execute_workflow cleanup", + workflow_run_id=workflow_run_id, + exc_info=True, + ) + if browser_cleanup_result is not None: + try: + async with asyncio.timeout(BROWSER_SESSION_WRITE_BACK_TIMEOUT): + await self._persist_workflow_browser_session_if_needed( + workflow=workflow, + workflow_run=workflow_run, + browser_state=browser_cleanup_result.browser_state, + close_browser_on_completion=browser_cleanup_result.close_browser_on_completion, + workflow_run_status=WorkflowRunStatus.completed, + ) + # Only suppress clean_up_workflow's write-back once this one has + # actually persisted; a failed/timed-out attempt must still be retried. + browser_cleanup_result.browser_session_write_back_attempted = True + except BaseException: + LOG.warning( + "Pre-finalization browser session write-back failed during execute_workflow cleanup", + workflow_run_id=workflow_run_id, + exc_info=True, + ) try: workflow_run = await asyncio.shield( self._finalize_workflow_run_status( @@ -2368,6 +2654,7 @@ class WorkflowService: browser_session_id=browser_session_id, close_browser_on_completion=close_browser_on_completion, need_call_webhook=need_call_webhook, + browser_cleanup_result=browser_cleanup_result, ) return workflow_run @@ -3971,25 +4258,16 @@ class WorkflowService: workflow_run_id=workflow_run_id, ) if proxy_session_id: - updated_headers = app.AGENT_FUNCTION.merge_proxy_session_extra_http_headers( - headers, proxy_session_id - ) - # The OSS AgentFunction stub returns the original headers; only persist when cloud injected a pin. - if updated_headers == headers: - return proxy_location_update: ProxyLocationInput | None = None if workflow_run.proxy_location is None: proxy_location_update = proxy_location or ProxyLocation.RESIDENTIAL_ISP - update_kwargs: dict[str, Any] = { - "workflow_run_id": workflow_run_id, - "extra_http_headers": updated_headers, - } - if proxy_location_update is not None: - update_kwargs["proxy_location"] = proxy_location_update - await app.DATABASE.workflow_runs.update_workflow_run(**update_kwargs) - workflow_run.extra_http_headers = updated_headers - if proxy_location_update is not None: - workflow_run.proxy_location = proxy_location_update + if not await self._persist_proxy_pin_headers( + workflow_run=workflow_run, + headers=headers, + proxy_session_id=proxy_session_id, + proxy_location_update=proxy_location_update, + ): + return LOG.info( "Applied LoginBlock credential proxy pin to workflow run", credential_id=credential_id, @@ -4005,6 +4283,30 @@ class WorkflowService: exc_info=True, ) + async def _persist_proxy_pin_headers( + self, + *, + workflow_run: WorkflowRun, + headers: dict[str, str], + proxy_session_id: str, + proxy_location_update: ProxyLocationInput | None = None, + ) -> bool: + updated_headers = app.AGENT_FUNCTION.merge_proxy_session_extra_http_headers(headers, proxy_session_id) + # The OSS AgentFunction stub returns the original headers; only persist when cloud injected a pin. + if updated_headers == headers: + return False + update_kwargs: dict[str, Any] = { + "workflow_run_id": workflow_run.workflow_run_id, + "extra_http_headers": updated_headers, + } + if proxy_location_update is not None: + update_kwargs["proxy_location"] = proxy_location_update + await app.DATABASE.workflow_runs.update_workflow_run(**update_kwargs) + workflow_run.extra_http_headers = updated_headers + if proxy_location_update is not None: + workflow_run.proxy_location = proxy_location_update + return True + async def _evaluate_debug_session_profile_decision( self, *, @@ -4414,6 +4716,7 @@ class WorkflowService: totp_verification_url: str | None = None, totp_identifier: str | None = None, persist_browser_session: bool = False, + pin_saved_session_ip: bool = False, browser_profile_id: str | None = None, browser_profile_key: str | None = None, model: dict[str, Any] | None = None, @@ -4449,6 +4752,7 @@ class WorkflowService: totp_verification_url=totp_verification_url, totp_identifier=totp_identifier, persist_browser_session=persist_browser_session, + pin_saved_session_ip=pin_saved_session_ip, browser_profile_id=browser_profile_id, browser_profile_key=browser_profile_key, model=model, @@ -4842,6 +5146,7 @@ class WorkflowService: totp_verification_url: str | None | object = _UNSET, totp_identifier: str | None | object = _UNSET, persist_browser_session: bool | None = None, + pin_saved_session_ip: bool | None = None, browser_profile_id: str | None | object = _UNSET, browser_profile_key: str | None | object = _UNSET, model: dict[str, Any] | None | object = _UNSET, @@ -4879,6 +5184,7 @@ class WorkflowService: totp_verification_url=totp_verification_url, totp_identifier=totp_identifier, persist_browser_session=persist_browser_session, + pin_saved_session_ip=pin_saved_session_ip, browser_profile_id=browser_profile_id, browser_profile_key=browser_profile_key, model=model, @@ -4909,6 +5215,7 @@ class WorkflowService: totp_verification_url=totp_verification_url, totp_identifier=totp_identifier, persist_browser_session=persist_browser_session, + pin_saved_session_ip=pin_saved_session_ip, browser_profile_id=browser_profile_id, browser_profile_key=browser_profile_key, model=model, @@ -5255,22 +5562,119 @@ class WorkflowService: }, ) if force_browser_session: + workflow_run = await app.DATABASE.workflow_runs.create_workflow_run( + workflow_permanent_id=workflow_permanent_id, + workflow_id=workflow_id, + organization_id=organization_id, + browser_session_id=workflow_request.browser_session_id, + browser_profile_id=browser_profile_id, + proxy_location=workflow_request.proxy_location, + webhook_callback_url=workflow_request.webhook_callback_url, + totp_verification_url=workflow_request.totp_verification_url, + totp_identifier=workflow_request.totp_identifier, + parent_workflow_run_id=parent_workflow_run_id, + max_screenshot_scrolling_times=workflow_request.max_screenshot_scrolls, + max_elapsed_time_minutes=max_elapsed_time_minutes, + 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, + debug_session_id=debug_session_id, + ai_fallback=workflow_request.ai_fallback, + code_gen=code_gen, + workflow_run_id=workflow_run_id, + trigger_type=trigger_type, + workflow_schedule_id=workflow_schedule_id, + ignore_inherited_workflow_system_prompt=ignore_inherited_workflow_system_prompt, + copilot_session_id=copilot_session_id, + ) LOG.info( "Force-creating browser session for workflow run", workflow_permanent_id=workflow_permanent_id, + workflow_run_id=workflow_run.workflow_run_id, organization_id=organization_id, ) - browser_session = await app.PERSISTENT_SESSIONS_MANAGER.create_session( - organization_id=organization_id, - proxy_location=workflow_request.proxy_location, - timeout_minutes=60, # 60 minutes default timeout for forced browser sessions - ) - browser_session_id = browser_session.persistent_browser_session_id - LOG.info( - "Browser session created for workflow run", - workflow_permanent_id=workflow_permanent_id, - browser_session_id=browser_session_id, - ) + forced_browser_profile_id = None + effective_proxy_location = workflow_request.proxy_location + pin_required = False + try: + workflow = await self.get_workflow(workflow_id=workflow_id) + # pin_required must be known before any awaited call that can raise, so the + # except path never falls through to creating an unprofiled pinned session. + effective_proxy_location = ( + workflow_request.proxy_location + if workflow_request.proxy_location is not None + else workflow.proxy_location + ) + pin_required = ( + workflow.persist_browser_session + and workflow.pin_saved_session_ip + and should_generate_proxy_session_id(effective_proxy_location) + ) + # Rotating credential profile keys need the persisted run id before rendering. + rotating_credential_selections = await self._select_rotating_credential_parameters_for_render( + workflow=workflow, + workflow_run=workflow_run, + organization_id=organization_id, + ) + forced_browser_profile_id = await self._resolve_managed_browser_profile_for_run_request( + workflow=workflow, + organization_id=organization_id, + workflow_request=workflow_request, + effective_proxy_location=effective_proxy_location, + extra_parameter_values=rotating_credential_selections, + ) + except Exception: + LOG.warning( + "Failed to resolve managed browser profile for forced browser session", + workflow_permanent_id=workflow_permanent_id, + workflow_id=workflow_id, + workflow_run_id=workflow_run.workflow_run_id, + organization_id=organization_id, + exc_info=True, + ) + if pin_required and forced_browser_profile_id is None: + LOG.info( + "Skipping forced browser session without managed browser profile for pinned workflow", + workflow_permanent_id=workflow_permanent_id, + workflow_id=workflow_id, + workflow_run_id=workflow_run.workflow_run_id, + organization_id=organization_id, + ) + browser_session = None + else: + try: + browser_session = await app.PERSISTENT_SESSIONS_MANAGER.create_session( + organization_id=organization_id, + proxy_location=workflow_request.proxy_location, + timeout_minutes=60, # 60 minutes default timeout for forced browser sessions + runnable_type=FORCED_WORKFLOW_SESSION_RUNNABLE_TYPE, + browser_profile_id=forced_browser_profile_id, + inherit_profile_proxy=True, + ) + browser_session_id = browser_session.persistent_browser_session_id + LOG.info( + "Browser session created for workflow run", + workflow_permanent_id=workflow_permanent_id, + workflow_run_id=workflow_run.workflow_run_id, + browser_session_id=browser_session_id, + ) + workflow_run = await app.DATABASE.workflow_runs.update_workflow_run( + workflow_run_id=workflow_run.workflow_run_id, + browser_session_id=browser_session_id, + ) + except Exception: + LOG.warning( + "Failed to force-create browser session for workflow run", + workflow_permanent_id=workflow_permanent_id, + workflow_id=workflow_id, + workflow_run_id=workflow_run.workflow_run_id, + organization_id=organization_id, + exc_info=True, + ) + + return workflow_run return await app.DATABASE.workflow_runs.create_workflow_run( workflow_permanent_id=workflow_permanent_id, @@ -6325,16 +6729,12 @@ class WorkflowService: errors=errors, ) - async def clean_up_workflow( + async def _clean_up_workflow_browser( self, - workflow: Workflow, workflow_run: WorkflowRun, - api_key: str | None = None, close_browser_on_completion: bool = True, - need_call_webhook: bool = True, browser_session_id: str | None = None, - ) -> None: - analytics.capture("skyvern-oss-agent-workflow-status", {"status": workflow_run.status}) + ) -> WorkflowBrowserCleanupResult: tasks = await self.get_tasks_by_workflow_run_id(workflow_run.workflow_run_id) # Look up child workflow runs (e.g. from task_v2 blocks) to flatten their @@ -6368,80 +6768,133 @@ class WorkflowService: organization_id=workflow_run.organization_id, child_workflow_run_ids=child_workflow_run_ids, ) + return WorkflowBrowserCleanupResult( + browser_state=browser_state, + tasks=tasks, + all_workflow_task_ids=all_workflow_task_ids, + child_workflow_run_ids=child_workflow_run_ids, + close_browser_on_completion=close_browser_on_completion, + ) + + async def _persist_workflow_browser_session_if_needed( + self, + workflow: Workflow, + workflow_run: WorkflowRun, + browser_state: BrowserState | None, + close_browser_on_completion: bool, + workflow_run_status: WorkflowRunStatus | None = None, + ) -> None: + if not ( + browser_state and workflow.persist_browser_session and browser_state.browser_artifacts.browser_session_dir + ): + return + + effective_workflow_run_status = workflow_run_status or workflow_run.status + if effective_workflow_run_status != WorkflowRunStatus.completed: + LOG.info( + "Skipped persisting browser session for non-completed workflow run", + workflow_run_id=workflow_run.workflow_run_id, + workflow_run_status=effective_workflow_run_status, + ) + return + + browser_profile = None + if workflow_run.browser_profile_id: + browser_profile = await app.DATABASE.browser_sessions.get_browser_profile( + workflow_run.browser_profile_id, + organization_id=workflow_run.organization_id, + ) + + if ( + browser_profile + and browser_profile.is_managed + # Only write back to a managed profile this workflow owns; a foreign + # managed id supplied explicitly is treated like a user profile. + and browser_profile.workflow_permanent_id == workflow_run.workflow_permanent_id + ): + if not close_browser_on_completion: + await persist_session_cookies( + browser_state.browser_context, + browser_state.browser_artifacts.browser_session_dir, + ) + await app.STORAGE.store_browser_profile( + workflow_run.organization_id, + profile_id=browser_profile.browser_profile_id, + directory=browser_state.browser_artifacts.browser_session_dir, + ) + LOG.info( + "Persisted managed browser profile for workflow run", + workflow_run_id=workflow_run.workflow_run_id, + browser_profile_id=browser_profile.browser_profile_id, + ) + elif browser_profile is None: + # No managed profile to write back to: either none was stamped, or a + # stamped managed profile was deleted between setup and finalization. + # Persist to the legacy session archive so the run's state isn't lost — + # the next run reseeds a managed profile from it. Only the deleted-mid-run + # case is unexpected, so warn there. + if workflow_run.browser_profile_id: + LOG.warning( + "Managed browser profile missing at finalization; persisting to legacy session archive instead", + workflow_run_id=workflow_run.workflow_run_id, + browser_profile_id=workflow_run.browser_profile_id, + ) + browser_session_storage_key = await self.get_workflow_browser_session_storage_key( + workflow=workflow, + workflow_run=workflow_run, + ) + if not close_browser_on_completion: + await persist_session_cookies( + browser_state.browser_context, + browser_state.browser_artifacts.browser_session_dir, + ) + await app.STORAGE.store_browser_session( + workflow_run.organization_id, + browser_session_storage_key, + browser_state.browser_artifacts.browser_session_dir, + ) + LOG.info( + "Persisted browser session for workflow run", + workflow_run_id=workflow_run.workflow_run_id, + browser_session_storage_key=browser_session_storage_key, + ) + + async def clean_up_workflow( + self, + workflow: Workflow, + workflow_run: WorkflowRun, + api_key: str | None = None, + close_browser_on_completion: bool = True, + need_call_webhook: bool = True, + browser_session_id: str | None = None, + browser_cleanup_result: WorkflowBrowserCleanupResult | None = None, + ) -> None: + analytics.capture("skyvern-oss-agent-workflow-status", {"status": workflow_run.status}) + if browser_cleanup_result is None: + browser_cleanup_result = await self._clean_up_workflow_browser( + workflow_run=workflow_run, + close_browser_on_completion=close_browser_on_completion, + browser_session_id=browser_session_id, + ) + + browser_state = browser_cleanup_result.browser_state + tasks = browser_cleanup_result.tasks + all_workflow_task_ids = browser_cleanup_result.all_workflow_task_ids + child_workflow_run_ids = browser_cleanup_result.child_workflow_run_ids + close_browser_on_completion = browser_cleanup_result.close_browser_on_completion if browser_state: await self.persist_video_data( browser_state, workflow, workflow_run, close_browser_on_completion=close_browser_on_completion ) if tasks: await self.persist_debug_artifacts(browser_state, tasks[-1], workflow, workflow_run) - if workflow.persist_browser_session and browser_state.browser_artifacts.browser_session_dir: - if workflow_run.status != WorkflowRunStatus.completed: - LOG.info( - "Skipped persisting browser session for non-completed workflow run", - workflow_run_id=workflow_run.workflow_run_id, - workflow_run_status=workflow_run.status, - ) - else: - browser_profile = None - if workflow_run.browser_profile_id: - browser_profile = await app.DATABASE.browser_sessions.get_browser_profile( - workflow_run.browser_profile_id, - organization_id=workflow_run.organization_id, - ) - - if ( - browser_profile - and browser_profile.is_managed - # Only write back to a managed profile this workflow owns; a foreign - # managed id supplied explicitly is treated like a user profile. - and browser_profile.workflow_permanent_id == workflow_run.workflow_permanent_id - ): - if not close_browser_on_completion: - await persist_session_cookies( - browser_state.browser_context, - browser_state.browser_artifacts.browser_session_dir, - ) - await app.STORAGE.store_browser_profile( - workflow_run.organization_id, - profile_id=browser_profile.browser_profile_id, - directory=browser_state.browser_artifacts.browser_session_dir, - ) - LOG.info( - "Persisted managed browser profile for workflow run", - workflow_run_id=workflow_run.workflow_run_id, - browser_profile_id=browser_profile.browser_profile_id, - ) - elif browser_profile is None: - # No managed profile to write back to: either none was stamped, or a - # stamped managed profile was deleted between setup and finalization. - # Persist to the legacy session archive so the run's state isn't lost — - # the next run reseeds a managed profile from it. Only the deleted-mid-run - # case is unexpected, so warn there. - if workflow_run.browser_profile_id: - LOG.warning( - "Managed browser profile missing at finalization; persisting to legacy session archive instead", - workflow_run_id=workflow_run.workflow_run_id, - browser_profile_id=workflow_run.browser_profile_id, - ) - browser_session_storage_key = await self.get_workflow_browser_session_storage_key( - workflow=workflow, - workflow_run=workflow_run, - ) - if not close_browser_on_completion: - await persist_session_cookies( - browser_state.browser_context, - browser_state.browser_artifacts.browser_session_dir, - ) - await app.STORAGE.store_browser_session( - workflow_run.organization_id, - browser_session_storage_key, - browser_state.browser_artifacts.browser_session_dir, - ) - LOG.info( - "Persisted browser session for workflow run", - workflow_run_id=workflow_run.workflow_run_id, - browser_session_storage_key=browser_session_storage_key, - ) + if not browser_cleanup_result.browser_session_write_back_attempted: + await self._persist_workflow_browser_session_if_needed( + workflow=workflow, + workflow_run=workflow_run, + browser_state=browser_state, + close_browser_on_completion=close_browser_on_completion, + ) await app.ARTIFACT_MANAGER.wait_for_upload_aiotasks(all_workflow_task_ids) @@ -6946,6 +7399,7 @@ class WorkflowService: totp_verification_url=runtime_workflow.totp_verification_url, totp_identifier=runtime_workflow.totp_identifier, persist_browser_session=runtime_workflow.persist_browser_session, + pin_saved_session_ip=runtime_workflow.pin_saved_session_ip, browser_profile_id=runtime_workflow.browser_profile_id, browser_profile_key=runtime_workflow.browser_profile_key, model=runtime_workflow.model, @@ -7085,6 +7539,9 @@ class WorkflowService: totp_verification_url=request.totp_verification_url, totp_identifier=request.totp_identifier, persist_browser_session=request.persist_browser_session, + pin_saved_session_ip=request.pin_saved_session_ip + if "pin_saved_session_ip" in request.model_fields_set + else existing_latest_workflow.pin_saved_session_ip, browser_profile_id=request.browser_profile_id, browser_profile_key=request.browser_profile_key, model=request.model, @@ -7129,6 +7586,7 @@ class WorkflowService: totp_verification_url=request.totp_verification_url, totp_identifier=request.totp_identifier, persist_browser_session=request.persist_browser_session, + pin_saved_session_ip=request.pin_saved_session_ip, browser_profile_id=request.browser_profile_id, browser_profile_key=request.browser_profile_key, model=request.model, diff --git a/skyvern/schemas/proxy_pinning.py b/skyvern/schemas/proxy_pinning.py index e1e486136..1c4e2178e 100644 --- a/skyvern/schemas/proxy_pinning.py +++ b/skyvern/schemas/proxy_pinning.py @@ -2,6 +2,7 @@ from __future__ import annotations import json import secrets +from hashlib import sha256 from typing import Any from skyvern.schemas.proxy_location import GeoTarget, ProxyLocation, ProxyLocationInput @@ -31,6 +32,12 @@ def generate_proxy_session_id(source_id: str) -> str: return proxy_session_id +def derive_proxy_session_id(*parts: str) -> str: + if not parts or any(not part or not part.strip() for part in parts): + raise ValueError("Cannot derive proxy session id from empty parts") + return sha256(":".join(parts).encode("utf-8")).hexdigest()[:10] + + def redact_proxy_session_id(value: str | None) -> str | None: if not value: return None diff --git a/skyvern/schemas/workflows.py b/skyvern/schemas/workflows.py index 3be51ba17..84e709858 100644 --- a/skyvern/schemas/workflows.py +++ b/skyvern/schemas/workflows.py @@ -1315,6 +1315,7 @@ class WorkflowCreateYAMLRequest(BaseModel): totp_verification_url: str | None = None totp_identifier: str | None = None persist_browser_session: bool = False + pin_saved_session_ip: bool = False browser_profile_id: str | None = None browser_profile_key: str | None = None model: dict[str, Any] | None = None diff --git a/skyvern/webeye/default_persistent_sessions_manager.py b/skyvern/webeye/default_persistent_sessions_manager.py index faf91a19e..1dc1b506c 100644 --- a/skyvern/webeye/default_persistent_sessions_manager.py +++ b/skyvern/webeye/default_persistent_sessions_manager.py @@ -349,6 +349,7 @@ class DefaultPersistentSessionsManager(PersistentSessionsManager): is_high_priority: bool = False, browser_profile_id: str | None = None, generate_browser_profile: bool = False, + inherit_profile_proxy: bool = False, wait_for_startup: bool = True, ) -> PersistentBrowserSession: """Create a new browser session for an organization and return its ID with the browser state.""" @@ -367,6 +368,7 @@ class DefaultPersistentSessionsManager(PersistentSessionsManager): browser_type=browser_type, browser_profile_id=browser_profile_id, generate_browser_profile=generate_browser_profile, + inherit_profile_proxy=inherit_profile_proxy, ) # Launch the browser immediately for standalone sessions so the diff --git a/skyvern/webeye/persistent_sessions_manager.py b/skyvern/webeye/persistent_sessions_manager.py index 9893a35a7..09d733a98 100644 --- a/skyvern/webeye/persistent_sessions_manager.py +++ b/skyvern/webeye/persistent_sessions_manager.py @@ -101,6 +101,7 @@ class PersistentSessionsManager(Protocol): is_high_priority: bool = False, browser_profile_id: str | None = None, generate_browser_profile: bool = False, + inherit_profile_proxy: bool = False, wait_for_startup: bool = True, ) -> PersistentBrowserSession: """Create a new browser session.""" diff --git a/tests/unit/forge/sdk/db/test_update_workflow_definition_fields.py b/tests/unit/forge/sdk/db/test_update_workflow_definition_fields.py index 3c8e91171..cadbe9696 100644 --- a/tests/unit/forge/sdk/db/test_update_workflow_definition_fields.py +++ b/tests/unit/forge/sdk/db/test_update_workflow_definition_fields.py @@ -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"} diff --git a/tests/unit/forge/sdk/db/test_workflow_runs_repository.py b/tests/unit/forge/sdk/db/test_workflow_runs_repository.py index 7759d45c2..293d71606 100644 --- a/tests/unit/forge/sdk/db/test_workflow_runs_repository.py +++ b/tests/unit/forge/sdk/db/test_workflow_runs_repository.py @@ -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 diff --git a/tests/unit/test_copilot_code_block_persist_seam.py b/tests/unit/test_copilot_code_block_persist_seam.py index 4d0d3a0c9..c242d4d96 100644 --- a/tests/unit/test_copilot_code_block_persist_seam.py +++ b/tests/unit/test_copilot_code_block_persist_seam.py @@ -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, diff --git a/tests/unit/test_credential_rotation.py b/tests/unit/test_credential_rotation.py index a52697405..8046a5f2b 100644 --- a/tests/unit/test_credential_rotation.py +++ b/tests/unit/test_credential_rotation.py @@ -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)", + ) diff --git a/tests/unit/test_login_block_credential_proxy_pin.py b/tests/unit/test_login_block_credential_proxy_pin.py index 1dc177c73..608b888a3 100644 --- a/tests/unit/test_login_block_credential_proxy_pin.py +++ b/tests/unit/test_login_block_credential_proxy_pin.py @@ -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) diff --git a/tests/unit/test_saved_session_proxy_pin.py b/tests/unit/test_saved_session_proxy_pin.py new file mode 100644 index 000000000..9ad4b87b1 --- /dev/null +++ b/tests/unit/test_saved_session_proxy_pin.py @@ -0,0 +1,1111 @@ +from datetime import UTC, datetime +from types import SimpleNamespace +from typing import Any +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from skyvern.forge import app +from skyvern.forge.sdk.db.repositories.browser_sessions import BrowserSessionsRepository +from skyvern.forge.sdk.schemas.browser_profiles import BrowserProfile +from skyvern.forge.sdk.schemas.persistent_browser_sessions import FORCED_WORKFLOW_SESSION_RUNNABLE_TYPE +from skyvern.forge.sdk.workflow.browser_profile_key import build_browser_profile_key_digest +from skyvern.forge.sdk.workflow.models.parameter import CredentialParameter, WorkflowParameter, WorkflowParameterType +from skyvern.forge.sdk.workflow.models.workflow import WorkflowRequestBody +from skyvern.forge.sdk.workflow.service import WorkflowService +from skyvern.schemas.proxy_pinning import derive_proxy_session_id, is_proxy_session_id +from skyvern.schemas.runs import ProxyLocation +from tests.unit.conftest import MockAsyncSessionCtx + + +def _workflow( + *, + persist_browser_session: bool = True, + pin_saved_session_ip: bool = True, + browser_profile_key: str | None = None, + proxy_location: ProxyLocation | None = ProxyLocation.RESIDENTIAL_ISP, + parameters: list[Any] | None = None, +) -> SimpleNamespace: + return SimpleNamespace( + persist_browser_session=persist_browser_session, + pin_saved_session_ip=pin_saved_session_ip, + browser_profile_key=browser_profile_key, + proxy_location=proxy_location, + workflow_permanent_id="wpid_test", + workflow_id="wf_test", + organization_id="org_test", + title="Workflow", + webhook_callback_url=None, + extra_http_headers=None, + cdp_connect_headers=None, + browser_profile_id=None, + max_elapsed_time_minutes=None, + run_with="agent", + code_version=None, + adaptive_caching=False, + sequential_key=None, + workflow_definition=SimpleNamespace(parameters=parameters or []), + ) + + +def _workflow_parameter(key: str, default_value: Any = None) -> WorkflowParameter: + now = datetime.now(UTC) + return WorkflowParameter( + workflow_parameter_id=f"wfp_{key}", + workflow_id="wf_test", + key=key, + workflow_parameter_type=WorkflowParameterType.STRING, + default_value=default_value, + created_at=now, + modified_at=now, + ) + + +def _credential_parameter( + key: str, + *, + credential_id: str = "cred_default", + credential_ids: list[str] | None = None, + selection_strategy: str | None = None, +) -> CredentialParameter: + now = datetime.now(UTC) + return CredentialParameter( + credential_parameter_id=f"cp_{key}", + workflow_id="wf_test", + key=key, + credential_id=credential_id, + credential_ids=credential_ids, + selection_strategy=selection_strategy, + created_at=now, + modified_at=now, + ) + + +def _workflow_run( + *, + proxy_location: ProxyLocation | str | None = ProxyLocation.RESIDENTIAL_ISP, + browser_profile_id: str | None = None, +) -> SimpleNamespace: + return SimpleNamespace( + workflow_run_id="wr_test", + workflow_permanent_id="wpid_test", + organization_id="org_test", + browser_profile_id=browser_profile_id, + proxy_location=proxy_location, + ) + + +def _profile( + *, + proxy_session_id: str | None = None, + is_managed: bool = True, + browser_profile_id: str = "bp_managed", +) -> BrowserProfile: + now = datetime.now(UTC) + return BrowserProfile( + browser_profile_id=browser_profile_id, + organization_id="org_test", + name="managed profile", + proxy_session_id=proxy_session_id, + proxy_location=ProxyLocation.RESIDENTIAL_ISP if proxy_session_id else None, + is_managed=is_managed, + created_at=now, + modified_at=now, + ) + + +def _mock_storage(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(app.STORAGE, "retrieve_browser_session", AsyncMock(return_value=None)) + monkeypatch.setattr(app.STORAGE, "store_browser_profile", AsyncMock()) + + +async def _create_forced_workflow_run( + monkeypatch: pytest.MonkeyPatch, + *, + workflow: SimpleNamespace, + workflow_request: WorkflowRequestBody | None = None, + get_or_create_profile: AsyncMock | None = None, + update_profile: AsyncMock | None = None, + create_session: AsyncMock | None = None, +) -> SimpleNamespace: + workflow_request = workflow_request or WorkflowRequestBody(proxy_location=ProxyLocation.RESIDENTIAL_ISP) + get_or_create_profile = get_or_create_profile or AsyncMock(return_value=(_profile(), False)) + update_profile = update_profile or AsyncMock(return_value=_profile()) + create_events: list[str] = [] + created_workflow_run = SimpleNamespace(workflow_run_id="wr_forced") + + async def _create_workflow_run(**_: object) -> SimpleNamespace: + create_events.append("create_workflow_run") + return created_workflow_run + + create_workflow_run = AsyncMock(side_effect=_create_workflow_run) + create_workflow_run.return_value = created_workflow_run + if create_session is None: + created_session = SimpleNamespace(persistent_browser_session_id="pbs_forced") + + async def _create_session(**_: object) -> SimpleNamespace: + create_events.append("create_session") + return created_session + + create_session = AsyncMock(side_effect=_create_session) + create_session.return_value = created_session + monkeypatch.setattr( + app.EXPERIMENTATION_PROVIDER, + "is_feature_enabled_cached", + AsyncMock(return_value=True), + ) + monkeypatch.setattr( + app.DATABASE.browser_sessions, + "get_or_create_managed_browser_profile", + get_or_create_profile, + ) + monkeypatch.setattr(app.DATABASE.browser_sessions, "update_browser_profile", update_profile) + monkeypatch.setattr(app.PERSISTENT_SESSIONS_MANAGER, "create_session", create_session) + monkeypatch.setattr(app.DATABASE.workflow_runs, "create_workflow_run", create_workflow_run) + update_workflow_run = AsyncMock( + return_value=SimpleNamespace(workflow_run_id="wr_forced", browser_session_id="pbs_forced") + ) + monkeypatch.setattr(app.DATABASE.workflow_runs, "update_workflow_run", update_workflow_run) + service = WorkflowService() + monkeypatch.setattr(service, "get_workflow", AsyncMock(return_value=workflow)) + + result = await service.create_workflow_run( + workflow_request=workflow_request, + workflow_permanent_id="wpid_test", + workflow_id="wf_test", + organization_id="org_test", + ) + + if "create_session" in create_events: + assert create_events.index("create_workflow_run") < create_events.index("create_session") + + return SimpleNamespace( + result=result, + create_session=create_session, + get_or_create_profile=get_or_create_profile, + update_profile=update_profile, + create_workflow_run=create_workflow_run, + update_workflow_run=update_workflow_run, + ) + + +async def _prepare_profile( + monkeypatch: pytest.MonkeyPatch, + *, + workflow: SimpleNamespace | None = None, + workflow_run: SimpleNamespace | None = None, + profile: SimpleNamespace | None = None, + update_profile: AsyncMock | None = None, + parameter_values: dict[str, object] | None = None, +) -> SimpleNamespace: + workflow = workflow or _workflow() + workflow_run = workflow_run or _workflow_run() + profile = profile or _profile() + update_profile = update_profile or AsyncMock(return_value=profile) + monkeypatch.setattr( + app.DATABASE.browser_sessions, + "get_or_create_managed_browser_profile", + AsyncMock(return_value=(profile, False)), + ) + monkeypatch.setattr(app.DATABASE.browser_sessions, "update_browser_profile", update_profile) + updated_run = SimpleNamespace(**{**vars(workflow_run), "browser_profile_id": profile.browser_profile_id}) + monkeypatch.setattr(app.DATABASE.workflow_runs, "update_workflow_run", AsyncMock(return_value=updated_run)) + _mock_storage(monkeypatch) + + return await WorkflowService()._prepare_persisted_workflow_browser_profile( + workflow=workflow, # type: ignore[arg-type] + workflow_run=workflow_run, # type: ignore[arg-type] + parameter_values=parameter_values or {"credential_id": "cred_a"}, + ) + + +async def _setup_profile_with_reconcile_failure( + monkeypatch: pytest.MonkeyPatch, + *, + workflow: SimpleNamespace, + workflow_request: WorkflowRequestBody | None = None, + profile: SimpleNamespace | None = None, +) -> tuple[WorkflowService, SimpleNamespace, AsyncMock, AsyncMock, Exception | None]: + workflow_request = workflow_request or WorkflowRequestBody() + profile = profile or _profile(proxy_session_id=None) + workflow_run = _workflow_run(proxy_location=workflow_request.proxy_location or workflow.proxy_location) + updated_run = SimpleNamespace(**{**vars(workflow_run), "browser_profile_id": profile.browser_profile_id}) + service = WorkflowService() + 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._select_rotating_credential_parameters_for_render = AsyncMock(return_value={}) # type: ignore[method-assign] + service._record_workflow_run_metadata_in_background = MagicMock() # type: ignore[method-assign] + mark_failed = AsyncMock(return_value=workflow_run) + service.mark_workflow_run_as_failed = mark_failed # type: ignore[method-assign] + update_profile = AsyncMock(side_effect=RuntimeError("db down")) + outer_session = AsyncMock() + + monkeypatch.setattr(app.DATABASE.workflow_runs, "Session", lambda: MockAsyncSessionCtx(outer_session)) + monkeypatch.setattr( + app.DATABASE.browser_sessions, + "get_or_create_managed_browser_profile", + AsyncMock(return_value=(profile, False)), + ) + monkeypatch.setattr(app.DATABASE.browser_sessions, "update_browser_profile", update_profile) + monkeypatch.setattr(app.DATABASE.workflow_runs, "update_workflow_run", AsyncMock(return_value=updated_run)) + monkeypatch.setattr("skyvern.forge.sdk.workflow.service.skyvern_context.current", lambda: None) + monkeypatch.setattr("skyvern.forge.sdk.workflow.service.skyvern_context.replace", MagicMock()) + _mock_storage(monkeypatch) + + caught: Exception | None = None + try: + result = await service.setup_workflow_run( + request_id="req_test", + workflow_request=workflow_request, + workflow_permanent_id="wpid_test", + organization=SimpleNamespace(organization_id="org_test", organization_name="Test Org"), + ) + except Exception as exc: + caught = exc + result = workflow_run + return service, result, update_profile, mark_failed, caught + + +def test_derive_proxy_session_id_supports_profile_segments() -> None: + proxy_session_id = derive_proxy_session_id("org_1", "wpid_1", "digest_a") + + assert len(proxy_session_id) == 10 + assert is_proxy_session_id(proxy_session_id) + assert derive_proxy_session_id("org_1", "wpid_1", "digest_a") == proxy_session_id + assert derive_proxy_session_id("org_1", "wpid_1", "digest_b") != proxy_session_id + + with pytest.raises(ValueError, match="empty parts"): + derive_proxy_session_id("org_1", "wpid_1", " ") + with pytest.raises(ValueError, match="empty parts"): + derive_proxy_session_id("org_1", None) # type: ignore[arg-type] + with pytest.raises(ValueError, match="empty parts"): + derive_proxy_session_id() + + +@pytest.mark.asyncio +async def test_force_browser_session_passes_managed_profile_and_pins_proxy( + monkeypatch: pytest.MonkeyPatch, +) -> None: + update_profile = AsyncMock(return_value=_profile()) + workflow = _workflow( + browser_profile_key="{{ credential_id }}", + parameters=[_workflow_parameter("credential_id", default_value="cred_default")], + ) + request = WorkflowRequestBody( + data={"credential_id": "cred_request"}, + proxy_location=ProxyLocation.RESIDENTIAL_ISP, + ) + + forced = await _create_forced_workflow_run( + monkeypatch, + workflow=workflow, + workflow_request=request, + update_profile=update_profile, + ) + + digest = build_browser_profile_key_digest("cred_request") + expected_pin = derive_proxy_session_id("org_test", "wpid_test", digest) + forced.get_or_create_profile.assert_awaited_once_with( + organization_id="org_test", + workflow_permanent_id="wpid_test", + browser_profile_key_digest=digest, + name="Workflow (auto-saved: cred_request)", + ) + update_profile.assert_awaited_once_with( + profile_id="bp_managed", + organization_id="org_test", + proxy_location=ProxyLocation.RESIDENTIAL_ISP, + proxy_session_id=expected_pin, + ) + forced.create_session.assert_awaited_once_with( + organization_id="org_test", + proxy_location=ProxyLocation.RESIDENTIAL_ISP, + timeout_minutes=60, + runnable_type=FORCED_WORKFLOW_SESSION_RUNNABLE_TYPE, + browser_profile_id="bp_managed", + inherit_profile_proxy=True, + ) + forced.create_workflow_run.assert_awaited_once() + assert forced.create_workflow_run.await_args.kwargs["browser_session_id"] is None + forced.update_workflow_run.assert_awaited_once_with( + workflow_run_id="wr_forced", + browser_session_id="pbs_forced", + ) + assert forced.result is forced.update_workflow_run.return_value + + +@pytest.mark.asyncio +async def test_force_browser_session_rotating_profile_key_selects_after_run_create( + monkeypatch: pytest.MonkeyPatch, +) -> None: + select_credential_for_run = AsyncMock(return_value="cred_selected") + monkeypatch.setattr( + "skyvern.forge.sdk.workflow.service.select_credential_for_run", + select_credential_for_run, + ) + update_profile = AsyncMock(return_value=_profile()) + workflow = _workflow( + browser_profile_key="{{ login_cred }}", + parameters=[ + _credential_parameter( + "login_cred", + credential_ids=["cred_a", "cred_selected"], + selection_strategy="round_robin", + ) + ], + ) + + forced = await _create_forced_workflow_run( + monkeypatch, + workflow=workflow, + workflow_request=WorkflowRequestBody( + data={"login_cred": "cred_request"}, + proxy_location=ProxyLocation.RESIDENTIAL_ISP, + ), + update_profile=update_profile, + ) + + digest = build_browser_profile_key_digest("cred_selected") + select_credential_for_run.assert_awaited_once_with( + workflow_run_id="wr_forced", + organization_id="org_test", + workflow_permanent_id="wpid_test", + parameter_key="login_cred", + credential_ids=["cred_a", "cred_selected"], + selection_strategy="round_robin", + ) + forced.get_or_create_profile.assert_awaited_once_with( + organization_id="org_test", + workflow_permanent_id="wpid_test", + browser_profile_key_digest=digest, + name="Workflow (auto-saved: cred_selected)", + ) + forced.create_session.assert_awaited_once_with( + organization_id="org_test", + proxy_location=ProxyLocation.RESIDENTIAL_ISP, + timeout_minutes=60, + runnable_type=FORCED_WORKFLOW_SESSION_RUNNABLE_TYPE, + browser_profile_id="bp_managed", + inherit_profile_proxy=True, + ) + forced.create_workflow_run.assert_awaited_once() + assert forced.create_workflow_run.await_args.kwargs["browser_session_id"] is None + forced.update_workflow_run.assert_awaited_once_with( + workflow_run_id="wr_forced", + browser_session_id="pbs_forced", + ) + assert forced.result is forced.update_workflow_run.return_value + + +@pytest.mark.asyncio +async def test_force_browser_session_rotating_profile_key_selection_failure_returns_run( + monkeypatch: pytest.MonkeyPatch, +) -> None: + select_credential_for_run = AsyncMock(side_effect=RuntimeError("selection failed")) + monkeypatch.setattr( + "skyvern.forge.sdk.workflow.service.select_credential_for_run", + select_credential_for_run, + ) + workflow = _workflow( + browser_profile_key="{{ login_cred }}", + parameters=[_credential_parameter("login_cred", credential_ids=["cred_a", "cred_b"])], + ) + + forced = await _create_forced_workflow_run( + monkeypatch, + workflow=workflow, + workflow_request=WorkflowRequestBody(proxy_location=ProxyLocation.RESIDENTIAL_ISP), + ) + + assert forced.result.workflow_run_id == "wr_forced" + select_credential_for_run.assert_awaited_once() + forced.create_workflow_run.assert_awaited_once() + assert forced.create_workflow_run.await_args.kwargs["browser_session_id"] is None + forced.get_or_create_profile.assert_not_awaited() + forced.create_session.assert_not_awaited() + forced.update_workflow_run.assert_not_awaited() + assert forced.result is forced.create_workflow_run.return_value + + +@pytest.mark.asyncio +async def test_force_browser_session_created_profile_seeds_legacy_session( + monkeypatch: pytest.MonkeyPatch, +) -> None: + profile = _profile() + get_or_create_profile = AsyncMock(return_value=(profile, True)) + retrieve_browser_session = AsyncMock(return_value="/tmp/legacy-session") + store_browser_profile = AsyncMock() + monkeypatch.setattr(app.STORAGE, "retrieve_browser_session", retrieve_browser_session) + monkeypatch.setattr(app.STORAGE, "store_browser_profile", store_browser_profile) + + forced = await _create_forced_workflow_run( + monkeypatch, + workflow=_workflow(), + get_or_create_profile=get_or_create_profile, + update_profile=AsyncMock(return_value=profile), + ) + + retrieve_browser_session.assert_awaited_once_with("org_test", "wpid_test") + store_browser_profile.assert_awaited_once_with( + "org_test", + profile_id="bp_managed", + directory="/tmp/legacy-session", + ) + forced.create_session.assert_awaited_once_with( + organization_id="org_test", + proxy_location=ProxyLocation.RESIDENTIAL_ISP, + timeout_minutes=60, + runnable_type=FORCED_WORKFLOW_SESSION_RUNNABLE_TYPE, + browser_profile_id="bp_managed", + inherit_profile_proxy=True, + ) + forced.create_workflow_run.assert_awaited_once() + assert forced.create_workflow_run.await_args.kwargs["browser_session_id"] is None + forced.update_workflow_run.assert_awaited_once_with( + workflow_run_id="wr_forced", + browser_session_id="pbs_forced", + ) + assert forced.result is forced.update_workflow_run.return_value + + +@pytest.mark.asyncio +async def test_force_browser_session_created_profile_seed_failure_rolls_back_and_skips_session( + monkeypatch: pytest.MonkeyPatch, +) -> None: + profile = _profile() + get_or_create_profile = AsyncMock(return_value=(profile, True)) + hard_delete = AsyncMock() + monkeypatch.setattr(app.STORAGE, "retrieve_browser_session", AsyncMock(return_value="/tmp/legacy-session")) + monkeypatch.setattr(app.STORAGE, "store_browser_profile", AsyncMock(side_effect=RuntimeError("storage down"))) + monkeypatch.setattr(app.DATABASE.browser_sessions, "hard_delete_browser_profile", hard_delete) + + forced = await _create_forced_workflow_run( + monkeypatch, + workflow=_workflow(), + get_or_create_profile=get_or_create_profile, + ) + + hard_delete.assert_awaited_once_with(profile_id="bp_managed", organization_id="org_test") + forced.update_profile.assert_not_awaited() + forced.create_session.assert_not_awaited() + forced.create_workflow_run.assert_awaited_once() + assert forced.create_workflow_run.await_args.kwargs["browser_session_id"] is None + forced.update_workflow_run.assert_not_awaited() + assert forced.result is forced.create_workflow_run.return_value + + +@pytest.mark.asyncio +async def test_force_browser_session_existing_profile_does_not_seed_legacy_session( + monkeypatch: pytest.MonkeyPatch, +) -> None: + retrieve_browser_session = AsyncMock() + store_browser_profile = AsyncMock() + monkeypatch.setattr(app.STORAGE, "retrieve_browser_session", retrieve_browser_session) + monkeypatch.setattr(app.STORAGE, "store_browser_profile", store_browser_profile) + + forced = await _create_forced_workflow_run( + monkeypatch, + workflow=_workflow(), + get_or_create_profile=AsyncMock(return_value=(_profile(), False)), + ) + + retrieve_browser_session.assert_not_awaited() + store_browser_profile.assert_not_awaited() + forced.create_workflow_run.assert_awaited_once() + assert forced.create_workflow_run.await_args.kwargs["browser_session_id"] is None + forced.update_workflow_run.assert_awaited_once_with( + workflow_run_id="wr_forced", + browser_session_id="pbs_forced", + ) + assert forced.result is forced.update_workflow_run.return_value + + +@pytest.mark.asyncio +async def test_force_browser_session_uses_workflow_proxy_default_for_profile_reconcile( + monkeypatch: pytest.MonkeyPatch, +) -> None: + update_profile = AsyncMock(return_value=_profile()) + + forced = await _create_forced_workflow_run( + monkeypatch, + workflow=_workflow(proxy_location=ProxyLocation.RESIDENTIAL_ISP), + workflow_request=WorkflowRequestBody(), + update_profile=update_profile, + ) + + expected_pin = derive_proxy_session_id("org_test", "wpid_test") + update_profile.assert_awaited_once_with( + profile_id="bp_managed", + organization_id="org_test", + proxy_location=ProxyLocation.RESIDENTIAL_ISP, + proxy_session_id=expected_pin, + ) + forced.create_workflow_run.assert_awaited_once() + assert forced.create_workflow_run.await_args.kwargs["browser_session_id"] is None + forced.update_workflow_run.assert_awaited_once_with( + workflow_run_id="wr_forced", + browser_session_id="pbs_forced", + ) + assert forced.result is forced.update_workflow_run.return_value + + +@pytest.mark.asyncio +async def test_force_browser_session_persist_off_does_not_pass_browser_profile( + monkeypatch: pytest.MonkeyPatch, +) -> None: + forced = await _create_forced_workflow_run( + monkeypatch, + workflow=_workflow(persist_browser_session=False), + ) + + forced.get_or_create_profile.assert_not_awaited() + forced.update_profile.assert_not_awaited() + forced.create_session.assert_awaited_once_with( + organization_id="org_test", + proxy_location=ProxyLocation.RESIDENTIAL_ISP, + timeout_minutes=60, + runnable_type=FORCED_WORKFLOW_SESSION_RUNNABLE_TYPE, + browser_profile_id=None, + inherit_profile_proxy=True, + ) + forced.update_workflow_run.assert_awaited_once_with( + workflow_run_id="wr_forced", + browser_session_id="pbs_forced", + ) + forced.create_workflow_run.assert_awaited_once() + assert forced.create_workflow_run.await_args.kwargs["browser_session_id"] is None + assert forced.result is forced.update_workflow_run.return_value + + +@pytest.mark.asyncio +async def test_force_browser_session_pinned_unresolvable_profile_key_skips_forced_session( + monkeypatch: pytest.MonkeyPatch, +) -> None: + forced = await _create_forced_workflow_run( + monkeypatch, + workflow=_workflow(browser_profile_key="{{ credential_id }}"), + workflow_request=WorkflowRequestBody(data={}, proxy_location=ProxyLocation.RESIDENTIAL_ISP), + ) + + forced.get_or_create_profile.assert_not_awaited() + forced.update_profile.assert_not_awaited() + forced.create_session.assert_not_awaited() + forced.create_workflow_run.assert_awaited_once() + assert forced.create_workflow_run.await_args.kwargs["browser_session_id"] is None + forced.update_workflow_run.assert_not_awaited() + assert forced.result is forced.create_workflow_run.return_value + + +@pytest.mark.asyncio +async def test_force_browser_session_pinned_profile_resolution_failure_skips_forced_session( + monkeypatch: pytest.MonkeyPatch, +) -> None: + forced = await _create_forced_workflow_run( + monkeypatch, + workflow=_workflow(), + get_or_create_profile=AsyncMock(side_effect=RuntimeError("db down")), + ) + + forced.get_or_create_profile.assert_awaited_once() + forced.update_profile.assert_not_awaited() + forced.create_session.assert_not_awaited() + forced.create_workflow_run.assert_awaited_once() + assert forced.create_workflow_run.await_args.kwargs["browser_session_id"] is None + forced.update_workflow_run.assert_not_awaited() + assert forced.result is forced.create_workflow_run.return_value + + +@pytest.mark.asyncio +async def test_force_browser_session_non_pinned_profile_resolution_failure_still_creates_session( + monkeypatch: pytest.MonkeyPatch, +) -> None: + forced = await _create_forced_workflow_run( + monkeypatch, + workflow=_workflow(pin_saved_session_ip=False), + get_or_create_profile=AsyncMock(side_effect=RuntimeError("db down")), + ) + + forced.get_or_create_profile.assert_awaited_once() + forced.update_profile.assert_not_awaited() + forced.create_session.assert_awaited_once_with( + organization_id="org_test", + proxy_location=ProxyLocation.RESIDENTIAL_ISP, + timeout_minutes=60, + runnable_type=FORCED_WORKFLOW_SESSION_RUNNABLE_TYPE, + browser_profile_id=None, + inherit_profile_proxy=True, + ) + forced.create_workflow_run.assert_awaited_once() + assert forced.create_workflow_run.await_args.kwargs["browser_session_id"] is None + forced.update_workflow_run.assert_awaited_once_with( + workflow_run_id="wr_forced", + browser_session_id="pbs_forced", + ) + assert forced.result is forced.update_workflow_run.return_value + + +@pytest.mark.asyncio +async def test_force_browser_session_non_pinned_unresolvable_profile_key_creates_unprofiled_session( + monkeypatch: pytest.MonkeyPatch, +) -> None: + forced = await _create_forced_workflow_run( + monkeypatch, + workflow=_workflow(browser_profile_key="{{ credential_id }}", pin_saved_session_ip=False), + workflow_request=WorkflowRequestBody(data={}, proxy_location=ProxyLocation.RESIDENTIAL_ISP), + ) + + forced.get_or_create_profile.assert_not_awaited() + forced.update_profile.assert_not_awaited() + forced.create_session.assert_awaited_once_with( + organization_id="org_test", + proxy_location=ProxyLocation.RESIDENTIAL_ISP, + timeout_minutes=60, + runnable_type=FORCED_WORKFLOW_SESSION_RUNNABLE_TYPE, + browser_profile_id=None, + inherit_profile_proxy=True, + ) + forced.create_workflow_run.assert_awaited_once() + assert forced.create_workflow_run.await_args.kwargs["browser_session_id"] is None + forced.update_workflow_run.assert_awaited_once_with( + workflow_run_id="wr_forced", + browser_session_id="pbs_forced", + ) + assert forced.result is forced.update_workflow_run.return_value + + +@pytest.mark.asyncio +async def test_force_browser_session_creation_failure_returns_run_without_session( + monkeypatch: pytest.MonkeyPatch, +) -> None: + forced = await _create_forced_workflow_run( + monkeypatch, + workflow=_workflow(), + create_session=AsyncMock(side_effect=RuntimeError("session creation failed")), + ) + + forced.create_workflow_run.assert_awaited_once() + assert forced.create_workflow_run.await_args.kwargs["browser_session_id"] is None + forced.create_session.assert_awaited_once() + forced.update_workflow_run.assert_not_awaited() + assert forced.result is forced.create_workflow_run.return_value + + +@pytest.mark.asyncio +async def test_create_workflow_run_non_force_path_single_create_no_update(monkeypatch: pytest.MonkeyPatch) -> None: + request = WorkflowRequestBody( + browser_session_id="pbs_requested", + proxy_location=ProxyLocation.RESIDENTIAL_ISP, + ) + create_workflow_run = AsyncMock(return_value=SimpleNamespace(workflow_run_id="wr_non_force")) + update_workflow_run = AsyncMock() + monkeypatch.setattr( + app.DATABASE.browser_sessions, + "get_persistent_browser_session", + AsyncMock(return_value=SimpleNamespace(browser_profile_id=None)), + ) + monkeypatch.setattr(app.DATABASE.workflow_runs, "create_workflow_run", create_workflow_run) + monkeypatch.setattr(app.DATABASE.workflow_runs, "update_workflow_run", update_workflow_run) + + result = await WorkflowService().create_workflow_run( + workflow_request=request, + workflow_permanent_id="wpid_test", + workflow_id="wf_test", + organization_id="org_test", + ) + + assert result.workflow_run_id == "wr_non_force" + create_workflow_run.assert_awaited_once_with( + workflow_permanent_id="wpid_test", + workflow_id="wf_test", + organization_id="org_test", + browser_session_id="pbs_requested", + browser_profile_id=None, + proxy_location=request.proxy_location, + webhook_callback_url=request.webhook_callback_url, + totp_verification_url=request.totp_verification_url, + totp_identifier=request.totp_identifier, + parent_workflow_run_id=None, + max_screenshot_scrolling_times=request.max_screenshot_scrolls, + max_elapsed_time_minutes=None, + extra_http_headers=request.extra_http_headers, + cdp_connect_headers=request.cdp_connect_headers, + browser_address=request.browser_address, + sequential_key=None, + run_with=request.run_with, + debug_session_id=None, + ai_fallback=request.ai_fallback, + code_gen=None, + workflow_run_id=None, + trigger_type=None, + workflow_schedule_id=None, + ignore_inherited_workflow_system_prompt=False, + copilot_session_id=None, + ) + update_workflow_run.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_prepare_managed_profile_sets_deterministic_pin( + monkeypatch: pytest.MonkeyPatch, +) -> None: + update_profile = AsyncMock(return_value=_profile()) + digest = build_browser_profile_key_digest("cred_a") + + result = await _prepare_profile( + monkeypatch, + workflow=_workflow(browser_profile_key="{{ credential_id }}"), + profile=_profile(proxy_session_id=None), + update_profile=update_profile, + ) + + expected_pin = derive_proxy_session_id("org_test", "wpid_test", digest) + assert result.browser_profile_id == "bp_managed" + update_profile.assert_awaited_once_with( + profile_id="bp_managed", + organization_id="org_test", + proxy_location=ProxyLocation.RESIDENTIAL_ISP, + proxy_session_id=expected_pin, + ) + + +@pytest.mark.asyncio +async def test_prepare_managed_profile_pins_keyless_workflow(monkeypatch: pytest.MonkeyPatch) -> None: + update_profile = AsyncMock(return_value=_profile()) + + await _prepare_profile( + monkeypatch, + workflow=_workflow(browser_profile_key=None), + profile=_profile(proxy_session_id=None), + update_profile=update_profile, + ) + + update_profile.assert_awaited_once_with( + profile_id="bp_managed", + organization_id="org_test", + proxy_location=ProxyLocation.RESIDENTIAL_ISP, + proxy_session_id=derive_proxy_session_id("org_test", "wpid_test"), + ) + + +@pytest.mark.asyncio +async def test_prepare_managed_profile_keeps_correct_existing_pin(monkeypatch: pytest.MonkeyPatch) -> None: + update_profile = AsyncMock() + expected_pin = derive_proxy_session_id("org_test", "wpid_test", build_browser_profile_key_digest("cred_a")) + + await _prepare_profile( + monkeypatch, + workflow=_workflow(browser_profile_key="{{ credential_id }}"), + profile=_profile(proxy_session_id=expected_pin), + update_profile=update_profile, + ) + + update_profile.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_prepare_managed_profile_heals_stale_proxy_location(monkeypatch: pytest.MonkeyPatch) -> None: + update_profile = AsyncMock(return_value=_profile()) + expected_pin = derive_proxy_session_id("org_test", "wpid_test", build_browser_profile_key_digest("cred_a")) + profile = _profile(proxy_session_id=expected_pin) + profile.proxy_location = None + + await _prepare_profile( + monkeypatch, + workflow=_workflow(browser_profile_key="{{ credential_id }}"), + profile=profile, + update_profile=update_profile, + ) + + update_profile.assert_awaited_once_with( + profile_id="bp_managed", + organization_id="org_test", + proxy_location=ProxyLocation.RESIDENTIAL_ISP, + proxy_session_id=expected_pin, + ) + + +@pytest.mark.asyncio +async def test_prepare_managed_profile_heals_drifted_pin(monkeypatch: pytest.MonkeyPatch) -> None: + update_profile = AsyncMock(return_value=_profile()) + expected_pin = derive_proxy_session_id("org_test", "wpid_test", build_browser_profile_key_digest("cred_a")) + + await _prepare_profile( + monkeypatch, + workflow=_workflow(browser_profile_key="{{ credential_id }}"), + profile=_profile(proxy_session_id="deadbeef99"), + update_profile=update_profile, + ) + + update_profile.assert_awaited_once_with( + profile_id="bp_managed", + organization_id="org_test", + proxy_location=ProxyLocation.RESIDENTIAL_ISP, + proxy_session_id=expected_pin, + ) + + +@pytest.mark.asyncio +async def test_prepare_managed_profile_clears_pin_when_toggle_off(monkeypatch: pytest.MonkeyPatch) -> None: + update_profile = AsyncMock(return_value=_profile(proxy_session_id=None)) + + await _prepare_profile( + monkeypatch, + workflow=_workflow(pin_saved_session_ip=False), + profile=_profile(proxy_session_id="abc1234567"), + update_profile=update_profile, + ) + + update_profile.assert_awaited_once_with( + profile_id="bp_managed", + organization_id="org_test", + proxy_location=None, + proxy_session_id=None, + ) + + +@pytest.mark.asyncio +async def test_prepare_managed_profile_clears_pin_for_non_isp_proxy_location(monkeypatch: pytest.MonkeyPatch) -> None: + update_profile = AsyncMock(return_value=_profile(proxy_session_id=None)) + + await _prepare_profile( + monkeypatch, + workflow_run=_workflow_run(proxy_location=ProxyLocation.RESIDENTIAL), + profile=_profile(proxy_session_id="abc1234567"), + update_profile=update_profile, + ) + + update_profile.assert_awaited_once_with( + profile_id="bp_managed", + organization_id="org_test", + proxy_location=None, + proxy_session_id=None, + ) + + +@pytest.mark.asyncio +async def test_prepare_managed_profile_does_not_set_pin_for_non_isp_proxy_location( + monkeypatch: pytest.MonkeyPatch, +) -> None: + update_profile = AsyncMock() + + await _prepare_profile( + monkeypatch, + workflow_run=_workflow_run(proxy_location=None), + profile=_profile(proxy_session_id=None), + update_profile=update_profile, + ) + + update_profile.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_setup_pinned_isp_managed_profile_reconcile_failure_fails_run( + monkeypatch: pytest.MonkeyPatch, +) -> None: + _, result, update_profile, mark_failed, caught = await _setup_profile_with_reconcile_failure( + monkeypatch, + workflow=_workflow(), + workflow_request=WorkflowRequestBody(proxy_location=ProxyLocation.RESIDENTIAL_ISP), + ) + + assert isinstance(caught, RuntimeError) + assert str(caught) == "db down" + assert result.browser_profile_id is None + update_profile.assert_awaited_once() + mark_failed.assert_awaited_once() + assert mark_failed.await_args.kwargs["workflow_run_id"] == "wr_test" + assert mark_failed.await_args.kwargs["failure_reason"].startswith("Setup workflow failed. failure reason:") + + +@pytest.mark.asyncio +async def test_setup_non_pinned_reconcile_failure_still_stamps_run(monkeypatch: pytest.MonkeyPatch) -> None: + _, result, update_profile, mark_failed, caught = await _setup_profile_with_reconcile_failure( + monkeypatch, + workflow=_workflow(pin_saved_session_ip=False), + profile=_profile(proxy_session_id="abc1234567"), + ) + + assert caught is None + assert result.browser_profile_id == "bp_managed" + update_profile.assert_awaited_once() + mark_failed.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_setup_pinned_non_isp_reconcile_failure_still_stamps_run(monkeypatch: pytest.MonkeyPatch) -> None: + _, result, update_profile, mark_failed, caught = await _setup_profile_with_reconcile_failure( + monkeypatch, + workflow=_workflow(proxy_location=ProxyLocation.RESIDENTIAL), + workflow_request=WorkflowRequestBody(proxy_location=ProxyLocation.RESIDENTIAL), + profile=_profile(proxy_session_id="abc1234567"), + ) + + assert caught is None + assert result.browser_profile_id == "bp_managed" + update_profile.assert_awaited_once() + mark_failed.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_prepare_user_profile_pin_is_never_cleared(monkeypatch: pytest.MonkeyPatch) -> None: + update_profile = AsyncMock() + + await _prepare_profile( + monkeypatch, + workflow=_workflow(pin_saved_session_ip=False), + profile=_profile(proxy_session_id="abc1234567", is_managed=False), + update_profile=update_profile, + ) + + update_profile.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_prepare_managed_profile_pin_is_deterministic_by_segment(monkeypatch: pytest.MonkeyPatch) -> None: + updates: list[str] = [] + + async def _record_update(**kwargs: object) -> SimpleNamespace: + updates.append(str(kwargs["proxy_session_id"])) + return _profile(proxy_session_id=str(kwargs["proxy_session_id"])) + + await _prepare_profile( + monkeypatch, + workflow=_workflow(browser_profile_key="{{ credential_id }}"), + profile=_profile(proxy_session_id=None, browser_profile_id="bp_first"), + update_profile=AsyncMock(side_effect=_record_update), + ) + await _prepare_profile( + monkeypatch, + workflow=_workflow(browser_profile_key="{{ credential_id }}"), + profile=_profile(proxy_session_id=None, browser_profile_id="bp_second"), + update_profile=AsyncMock(side_effect=_record_update), + ) + await _prepare_profile( + monkeypatch, + workflow=_workflow(browser_profile_key="{{ credential_id }}"), + profile=_profile(proxy_session_id=None, browser_profile_id="bp_third"), + update_profile=AsyncMock(side_effect=_record_update), + parameter_values={"credential_id": "cred_b"}, + ) + + same_segment_pin = derive_proxy_session_id("org_test", "wpid_test", build_browser_profile_key_digest("cred_a")) + different_segment_pin = derive_proxy_session_id( + "org_test", + "wpid_test", + build_browser_profile_key_digest("cred_b"), + ) + assert updates == [same_segment_pin, same_segment_pin, different_segment_pin] + assert different_segment_pin != same_segment_pin + + +class _SessionScalars: + def __init__(self, profile: SimpleNamespace | None) -> None: + self._profile = profile + + def first(self) -> SimpleNamespace | None: + return self._profile + + +async def _create_repo_session( + *, + profile: SimpleNamespace | None, + proxy_location: ProxyLocation | str | None = ProxyLocation.RESIDENTIAL_ISP, + proxy_session_id: str | None = None, + browser_profile_id: str | None = "bp_managed", + inherit_profile_proxy: bool = False, + expect_profile_lookup: bool | None = None, +) -> SimpleNamespace: + mock_session = AsyncMock() + mock_session.scalars.return_value = _SessionScalars(profile) + mock_session.add = MagicMock() + + async def _flush() -> None: + stored_session = mock_session.add.call_args.args[0] + if getattr(stored_session, "persistent_browser_session_id", None) is None: + stored_session.persistent_browser_session_id = "pbs_test" + + mock_session.flush = AsyncMock(side_effect=_flush) + mock_session.commit = AsyncMock() + mock_session.refresh = AsyncMock() + repo = BrowserSessionsRepository(session_factory=lambda: MockAsyncSessionCtx(mock_session)) + + def _validate(model: object) -> SimpleNamespace: + return SimpleNamespace( + persistent_browser_session_id=model.persistent_browser_session_id, + proxy_location=model.proxy_location, + proxy_session_id=model.proxy_session_id, + browser_profile_id=model.browser_profile_id, + ) + + with patch( + "skyvern.forge.sdk.schemas.persistent_browser_sessions.PersistentBrowserSession.model_validate", + side_effect=_validate, + ): + created_session = await repo.create_persistent_browser_session( + organization_id="org_test", + proxy_location=proxy_location, + proxy_session_id=proxy_session_id, + browser_profile_id=browser_profile_id, + inherit_profile_proxy=inherit_profile_proxy, + ) + if expect_profile_lookup is True: + mock_session.scalars.assert_awaited_once() + elif expect_profile_lookup is False: + mock_session.scalars.assert_not_awaited() + return created_session + + +@pytest.mark.asyncio +async def test_create_persistent_browser_session_inherits_profile_pin() -> None: + # The ORM row exposes proxy_location as a serialized string; the repo must deserialize it + # before it reaches serialize_proxy_location (which rejects a bare str). + profile = SimpleNamespace( + browser_profile_id="bp_managed", + proxy_session_id="abc1234567", + proxy_location=ProxyLocation.RESIDENTIAL_ISP.value, + is_managed=True, + ) + session = await _create_repo_session(profile=profile, inherit_profile_proxy=True) + + assert session.proxy_location == ProxyLocation.RESIDENTIAL_ISP.value + assert session.proxy_session_id == "abc1234567" + assert session.browser_profile_id == "bp_managed" + + +@pytest.mark.asyncio +async def test_create_persistent_browser_session_explicit_pin_wins() -> None: + session = await _create_repo_session( + profile=_profile(proxy_session_id="abc1234567"), + proxy_session_id="fff1234567", + inherit_profile_proxy=True, + ) + + assert session.proxy_location == ProxyLocation.RESIDENTIAL_ISP.value + assert session.proxy_session_id == "fff1234567" + + +@pytest.mark.asyncio +async def test_create_persistent_browser_session_pinless_profile_keeps_auto_generate_behavior() -> None: + session = await _create_repo_session(profile=_profile(proxy_session_id=None), inherit_profile_proxy=True) + + assert session.proxy_location == ProxyLocation.RESIDENTIAL_ISP.value + assert is_proxy_session_id(session.proxy_session_id) + assert session.proxy_session_id != "abc1234567" + + +@pytest.mark.asyncio +async def test_create_persistent_browser_session_does_not_inherit_profile_pin_by_default() -> None: + profile = _profile(proxy_session_id="abc1234567") + + session = await _create_repo_session( + profile=profile, + proxy_location=None, + browser_profile_id="bp_managed", + expect_profile_lookup=False, + ) + + assert session.proxy_location is None + assert session.proxy_session_id is None + assert session.browser_profile_id == "bp_managed" diff --git a/tests/unit/test_workflow_browser_profile_id.py b/tests/unit/test_workflow_browser_profile_id.py index c3d0dd44c..983f46bfa 100644 --- a/tests/unit/test_workflow_browser_profile_id.py +++ b/tests/unit/test_workflow_browser_profile_id.py @@ -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") diff --git a/tests/unit/test_workflow_elapsed_timeout.py b/tests/unit/test_workflow_elapsed_timeout.py index 775503f40..4b89d14d8 100644 --- a/tests/unit/test_workflow_elapsed_timeout.py +++ b/tests/unit/test_workflow_elapsed_timeout.py @@ -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", diff --git a/tests/unit/test_workflow_service_managed_browser_profile.py b/tests/unit/test_workflow_service_managed_browser_profile.py index a425c54b0..e91155a8e 100644 --- a/tests/unit/test_workflow_service_managed_browser_profile.py +++ b/tests/unit/test_workflow_service_managed_browser_profile.py @@ -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"]