From 869a94dbf44318cc83131fd62affc4ee1f5dd11c Mon Sep 17 00:00:00 2001
From: Shuchang Zheng <wintonzheng0325@gmail.com>
Date: Wed, 7 May 2025 13:04:02 -0700
Subject: [PATCH] support CUA navigation block (#2306)

---
 skyvern/forge/sdk/workflow/models/block.py | 3 +++
 skyvern/forge/sdk/workflow/models/yaml.py  | 3 ++-
 skyvern/forge/sdk/workflow/service.py      | 1 +
 3 files changed, 6 insertions(+), 1 deletion(-)

diff --git a/skyvern/forge/sdk/workflow/models/block.py b/skyvern/forge/sdk/workflow/models/block.py
index 63d25815..f3f52334 100644
--- a/skyvern/forge/sdk/workflow/models/block.py
+++ b/skyvern/forge/sdk/workflow/models/block.py
@@ -71,6 +71,7 @@ from skyvern.forge.sdk.workflow.models.parameter import (
     OutputParameter,
     WorkflowParameter,
 )
+from skyvern.schemas.runs import RunEngine
 from skyvern.utils.url_validators import prepend_scheme_and_validate_url
 from skyvern.webeye.browser_factory import BrowserState
 from skyvern.webeye.utils.page import SkyvernFrame
@@ -349,6 +350,7 @@ class BaseTaskBlock(Block):
     task_type: str = TaskType.general
     url: str | None = None
     title: str = ""
+    engine: RunEngine = RunEngine.skyvern_v1
     complete_criterion: str | None = None
     terminate_criterion: str | None = None
     navigation_goal: str | None = None
@@ -623,6 +625,7 @@ class BaseTaskBlock(Block):
                     browser_session_id=browser_session_id,
                     close_browser_on_completion=browser_session_id is None,
                     complete_verification=self.complete_verification,
+                    engine=self.engine,
                 )
             except Exception as e:
                 # Make sure the task is marked as failed in the database before raising the exception
diff --git a/skyvern/forge/sdk/workflow/models/yaml.py b/skyvern/forge/sdk/workflow/models/yaml.py
index 7212d74c..5fbc908e 100644
--- a/skyvern/forge/sdk/workflow/models/yaml.py
+++ b/skyvern/forge/sdk/workflow/models/yaml.py
@@ -8,7 +8,7 @@ from skyvern.forge.sdk.workflow.models.block import BlockType, FileType
 from skyvern.forge.sdk.workflow.models.constants import FileStorageType
 from skyvern.forge.sdk.workflow.models.parameter import ParameterType, WorkflowParameterType
 from skyvern.forge.sdk.workflow.models.workflow import WorkflowStatus
-from skyvern.schemas.runs import ProxyLocation
+from skyvern.schemas.runs import ProxyLocation, RunEngine
 
 
 class ParameterYAML(BaseModel, abc.ABC):
@@ -275,6 +275,7 @@ class NavigationBlockYAML(BlockYAML):
     navigation_goal: str
     url: str | None = None
     title: str = ""
+    engine: RunEngine = RunEngine.skyvern_v1
     error_code_mapping: dict[str, str] | None = None
     max_retries: int = 0
     max_steps_per_run: int | None = None
diff --git a/skyvern/forge/sdk/workflow/service.py b/skyvern/forge/sdk/workflow/service.py
index c0cb1d3f..19c15c71 100644
--- a/skyvern/forge/sdk/workflow/service.py
+++ b/skyvern/forge/sdk/workflow/service.py
@@ -1790,6 +1790,7 @@ class WorkflowService:
                 label=block_yaml.label,
                 url=block_yaml.url,
                 title=block_yaml.title,
+                engine=block_yaml.engine,
                 parameters=navigation_block_parameters,
                 output_parameter=output_parameter,
                 navigation_goal=block_yaml.navigation_goal,