CUA feature - verification code / 2FA (#2174)

skyvern/webeye/actions/actions.py

@@ -33,6 +33,7 @@ class ActionType(StrEnum):
     TYPE = "type"
     MOVE = "move"
     DRAG = "drag"
+    VERIFICATION_CODE = "verification_code"
 
     def is_web_action(self) -> bool:
         return self in [
@@ -293,6 +294,11 @@ class DragAction(Action):
     path: list[tuple[int, int]] = []
 
 
+class VerificationCodeAction(Action):
+    action_type: ActionType = ActionType.VERIFICATION_CODE
+    verification_code: str
+
+
 class ScrapeResult(BaseModel):
     """
     Scraped response from a webpage, including:

skyvern/webeye/actions/handler.py

@@ -1559,6 +1559,24 @@ async def handle_drag_action(
     return [ActionSuccess()]
 
 
+async def handle_verification_code_action(
+    action: actions.VerificationCodeAction,
+    page: Page,
+    scraped_page: ScrapedPage,
+    task: Task,
+    step: Step,
+) -> list[ActionResult]:
+    LOG.info(
+        "Setting verification code in skyvern context",
+        task_id=task.task_id,
+        step_id=step.step_id,
+        verification_code=action.verification_code,
+    )
+    current_context = skyvern_context.ensure_context()
+    current_context.totp_codes[task.task_id] = action.verification_code
+    return [ActionSuccess()]
+
+
 ActionHandler.register_action_type(ActionType.SOLVE_CAPTCHA, handle_solve_captcha_action)
 ActionHandler.register_action_type(ActionType.CLICK, handle_click_action)
 ActionHandler.register_action_type(ActionType.INPUT_TEXT, handle_input_text_action)
@@ -1574,6 +1592,7 @@ ActionHandler.register_action_type(ActionType.SCROLL, handle_scroll_action)
 ActionHandler.register_action_type(ActionType.KEYPRESS, handle_keypress_action)
 ActionHandler.register_action_type(ActionType.MOVE, handle_move_action)
 ActionHandler.register_action_type(ActionType.DRAG, handle_drag_action)
+ActionHandler.register_action_type(ActionType.VERIFICATION_CODE, handle_verification_code_action)
 
 
 async def get_actual_value_of_parameter_if_secret(task: Task, parameter: str) -> Any:

skyvern/webeye/actions/parse_actions.py

@@ -27,8 +27,10 @@ from skyvern.webeye.actions.actions import (
     SolveCaptchaAction,
     TerminateAction,
     UploadFileAction,
+    VerificationCodeAction,
     WaitAction,
 )
+from skyvern.webeye.actions.handler import poll_verification_code
 from skyvern.webeye.scraper.scraper import ScrapedPage
 
 LOG = structlog.get_logger()
@@ -400,12 +402,33 @@ async def parse_cua_actions(
                 intention=reasoning,
             )
         elif skyvern_action_type == "get_verification_code":
-            # Currently we don't support verification code
-            # TODO: handle verification code by fetching the code and send it to CUA
-            action = TerminateAction(
-                reasoning=reasoning,
-                intention=reasoning,
-            )
+            if (task.totp_verification_url or task.totp_identifier) and task.organization_id:
+                LOG.info(
+                    "Getting verification code for CUA",
+                    task_id=task.task_id,
+                    organization_id=task.organization_id,
+                    workflow_run_id=task.workflow_run_id,
+                    totp_verification_url=task.totp_verification_url,
+                    totp_identifier=task.totp_identifier,
+                )
+                verification_code = await poll_verification_code(
+                    task.task_id,
+                    task.organization_id,
+                    workflow_run_id=task.workflow_run_id,
+                    totp_verification_url=task.totp_verification_url,
+                    totp_identifier=task.totp_identifier,
+                )
+                reasoning = reasoning or f"Received verification code: {verification_code}"
+                action = VerificationCodeAction(
+                    verification_code=verification_code,
+                    reasoning=reasoning,
+                    intention=reasoning,
+                )
+            else:
+                action = TerminateAction(
+                    reasoning=reasoning,
+                    intention=reasoning,
+                )
 
         action.organization_id = task.organization_id
         action.workflow_run_id = task.workflow_run_id