From dde6837df87030471af10adfcd462973660a0fc5 Mon Sep 17 00:00:00 2001
From: Shuchang Zheng <wintonzheng0325@gmail.com>
Date: Fri, 12 Sep 2025 12:41:23 -0700
Subject: [PATCH] shu/fix totp script gen (#3423)

---
 skyvern/core/script_generations/generate_script.py | 4 ++--
 skyvern/core/totp.py                               | 8 ++++++++
 skyvern/forge/agent.py                             | 3 ++-
 skyvern/forge/sdk/core/skyvern_context.py          | 4 ++++
 4 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/skyvern/core/script_generations/generate_script.py b/skyvern/core/script_generations/generate_script.py
index 09d04370..de338975 100644
--- a/skyvern/core/script_generations/generate_script.py
+++ b/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),
diff --git a/skyvern/core/totp.py b/skyvern/core/totp.py
index 4820b2d2..ccd69bfc 100644
--- a/skyvern/core/totp.py
+++ b/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
diff --git a/skyvern/forge/agent.py b/skyvern/forge/agent.py
index 411ca393..1629f5c2 100644
--- a/skyvern/forge/agent.py
+++ b/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):
diff --git a/skyvern/forge/sdk/core/skyvern_context.py b/skyvern/forge/sdk/core/skyvern_context.py
index cc09240f..de117fcd 100644
--- a/skyvern/forge/sdk/core/skyvern_context.py
+++ b/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",