shu/fix totp script gen (#3423)

skyvern/core/script_generations/generate_script.py

@@ -334,7 +334,7 @@ def _action_to_stmt(act: dict[str, Any], task: dict[str, Any], assign_to_output:
                 args.append(
                     cst.Arg(
                         keyword=cst.Name("totp_identifier"),
-                        value=cst.Name(task.get("totp_identifier")),
+                        value=_value(task.get("totp_identifier")),
                         whitespace_after_arg=cst.ParenthesizedWhitespace(
                             indent=True,
                             last_line=cst.SimpleWhitespace(INDENT),
@@ -345,7 +345,7 @@ def _action_to_stmt(act: dict[str, Any], task: dict[str, Any], assign_to_output:
                 args.append(
                     cst.Arg(
                         keyword=cst.Name("totp_url"),
-                        value=cst.Name(task.get("totp_verification_url")),
+                        value=_value(task.get("totp_verification_url")),
                         whitespace_after_arg=cst.ParenthesizedWhitespace(
                             indent=True,
                             last_line=cst.SimpleWhitespace(INDENT),

skyvern/core/totp.py

@@ -30,6 +30,14 @@ async def poll_verification_code(
     if not org_token:
         LOG.error("Failed to get organization token when trying to get verification code")
         return None
+    LOG.info(
+        "Polling verification code",
+        task_id=task_id,
+        workflow_run_id=workflow_run_id,
+        workflow_permanent_id=workflow_permanent_id,
+        totp_verification_url=totp_verification_url,
+        totp_identifier=totp_identifier,
+    )
     while True:
         await asyncio.sleep(10)
         # check timeout

skyvern/forge/agent.py

@@ -999,6 +999,8 @@ class ForgeAgent:
                         )
                         detailed_agent_step_output.llm_response = json_response
                         actions = parse_actions(task, step.step_id, step.order, scraped_page, json_response["actions"])
+                        if context:
+                            context.pop_totp_code(task.task_id)
                     except NoTOTPVerificationCodeFound:
                         actions = [
                             TerminateAction(
@@ -3003,7 +3005,6 @@ class ForgeAgent:
                 browser_state,
                 scraped_page,
                 verification_code_check=False,
-                expire_verification_code=True,
             )
             llm_key_override = task.llm_key
             if await is_cua_task(task=task):

skyvern/forge/sdk/core/skyvern_context.py

@@ -39,6 +39,10 @@ class SkyvernContext:
     def __str__(self) -> str:
         return self.__repr__()
 
+    def pop_totp_code(self, task_id: str) -> None:
+        if task_id in self.totp_codes:
+            self.totp_codes.pop(task_id)
+
 
 _context: ContextVar[SkyvernContext | None] = ContextVar(
     "Global context",