From 2b365bce598cc33ae07c57deacc27e03d9ada9a0 Mon Sep 17 00:00:00 2001
From: Shuchang Zheng <wintonzheng0325@gmail.com>
Date: Wed, 11 Sep 2024 00:03:09 -0700
Subject: [PATCH] make totp code fetch more reliable (#809)

---
 skyvern/config.py                 | 1 +
 skyvern/forge/agent.py            | 2 +-
 skyvern/webeye/actions/handler.py | 2 ++
 3 files changed, 4 insertions(+), 1 deletion(-)

diff --git a/skyvern/config.py b/skyvern/config.py
index d2235dd2..ce2eb7da 100644
--- a/skyvern/config.py
+++ b/skyvern/config.py
@@ -119,6 +119,7 @@ class Settings(BaseSettings):
 
     # TOTP Settings
     TOTP_LIFESPAN_MINUTES: int = 10
+    VERIFICATION_CODE_INITIAL_WAIT_TIME_SECS: int = 40
     VERIFICATION_CODE_POLLING_TIMEOUT_MINS: int = 5
 
     def is_cloud_environment(self) -> bool:
diff --git a/skyvern/forge/agent.py b/skyvern/forge/agent.py
index 5fd5cfba..61252e22 100644
--- a/skyvern/forge/agent.py
+++ b/skyvern/forge/agent.py
@@ -1617,7 +1617,7 @@ class ForgeAgent:
                 browser_state,
                 element_tree_in_prompt,
                 verification_code_check=False,
-                expire_verification_code=False,
+                expire_verification_code=True,
             )
             return await app.LLM_API_HANDLER(
                 prompt=extract_action_prompt,
diff --git a/skyvern/webeye/actions/handler.py b/skyvern/webeye/actions/handler.py
index fa8ae073..d7fc4501 100644
--- a/skyvern/webeye/actions/handler.py
+++ b/skyvern/webeye/actions/handler.py
@@ -2118,6 +2118,8 @@ async def poll_verification_code(
     if not org_token:
         LOG.error("Failed to get organization token when trying to get verification code")
         return None
+    # wait for 40 seconds to let the verification code comes in before polling
+    await asyncio.sleep(SettingsManager.get_settings().VERIFICATION_CODE_INITIAL_WAIT_TIME_SECS)
     while True:
         # check timeout
         if datetime.utcnow() > timeout_datetime: