trigger TerminateAction when totp code is not found (#1768)

skyvern/exceptions.py

@@ -594,3 +594,23 @@ class UrlGenerationFailure(SkyvernHTTPException):
 class ObserverCruiseNotFound(SkyvernHTTPException):
     def __init__(self, observer_cruise_id: str) -> None:
         super().__init__(f"Observer task {observer_cruise_id} not found")
+
+
+class NoTOTPVerificationCodeFound(SkyvernHTTPException):
+    def __init__(
+        self,
+        task_id: str | None = None,
+        workflow_run_id: str | None = None,
+        totp_verification_url: str | None = None,
+        totp_identifier: str | None = None,
+    ) -> None:
+        msg = "No TOTP verification code found."
+        if task_id:
+            msg += f" task_id={task_id}"
+        if workflow_run_id:
+            msg += f" workflow_run_id={workflow_run_id}"
+        if totp_verification_url:
+            msg += f" totp_verification_url={totp_verification_url}"
+        if totp_identifier:
+            msg += f" totp_identifier={totp_identifier}"
+        super().__init__(msg)

skyvern/forge/agent.py

@@ -35,6 +35,7 @@ from skyvern.exceptions import (
     InvalidWorkflowTaskURLState,
     MissingBrowserState,
     MissingBrowserStatePage,
+    NoTOTPVerificationCodeFound,
     SkyvernException,
     StepTerminationError,
     StepUnableToExecuteError,
@@ -74,6 +75,7 @@ from skyvern.webeye.actions.actions import (
     DecisiveAction,
     ExtractAction,
     ReloadPageAction,
+    TerminateAction,
     UserDefinedError,
     WebAction,
 )
@@ -787,16 +789,29 @@ class ForgeAgent:
                     step=step,
                     screenshots=scraped_page.screenshots,
                 )
-                json_response = await self.handle_potential_verification_code(
+                try:
-                    task,
+                    json_response = await self.handle_potential_verification_code(
-                    step,
+                        task,
-                    scraped_page,
+                        step,
-                    browser_state,
+                        scraped_page,
-                    json_response,
+                        browser_state,
-                )
+                        json_response,
-                detailed_agent_step_output.llm_response = json_response
+                    )
-
+                    detailed_agent_step_output.llm_response = json_response
-                actions = parse_actions(task, step.step_id, step.order, scraped_page, json_response["actions"])
+                    actions = parse_actions(task, step.step_id, step.order, scraped_page, json_response["actions"])
+                except NoTOTPVerificationCodeFound:
+                    actions = [
+                        TerminateAction(
+                            organization_id=task.organization_id,
+                            workflow_run_id=task.workflow_run_id,
+                            task_id=task.task_id,
+                            step_id=step.step_id,
+                            step_order=step.order,
+                            action_order=0,
+                            reasoning="No TOTP verification code found. Going to terminate.",
+                            intention="No TOTP verification code found. Going to terminate.",
+                        )
+                    ]
 
             detailed_agent_step_output.actions = actions
             if len(actions) == 0:

skyvern/webeye/actions/handler.py

@@ -47,6 +47,7 @@ from skyvern.exceptions import (
     NoIncrementalElementFoundForAutoCompletion,
     NoIncrementalElementFoundForCustomSelection,
     NoSuitableAutoCompleteOption,
+    NoTOTPVerificationCodeFound,
     OptionIndexOutOfBound,
     WrongElementToUploadFile,
 )
@@ -2847,7 +2848,12 @@ async def poll_verification_code(
         # check timeout
         if datetime.utcnow() > timeout_datetime:
             LOG.warning("Polling verification code timed out", workflow_id=workflow_id)
-            return None
+            raise NoTOTPVerificationCodeFound(
+                task_id=task_id,
+                workflow_run_id=workflow_run_id,
+                totp_verification_url=totp_verification_url,
+                totp_identifier=totp_identifier,
+            )
         verification_code = None
         if totp_verification_url:
             verification_code = await _get_verification_code_from_url(task_id, totp_verification_url, org_token.token)