From b42a1900e2d361292617d6ec0dc6b68e473ee2f9 Mon Sep 17 00:00:00 2001
From: Shuchang Zheng <wintonzheng0325@gmail.com>
Date: Thu, 23 Jan 2025 22:04:32 +0800
Subject: [PATCH] optimize downalod waiting time logic (#1626)

---
 skyvern/constants.py              |  1 +
 skyvern/webeye/actions/handler.py | 35 +++++++++++++++----------------
 2 files changed, 18 insertions(+), 18 deletions(-)

diff --git a/skyvern/constants.py b/skyvern/constants.py
index fb644d7c..5e6455db 100644
--- a/skyvern/constants.py
+++ b/skyvern/constants.py
@@ -10,6 +10,7 @@ INPUT_TEXT_TIMEOUT = 120000  # 2 minutes
 PAGE_CONTENT_TIMEOUT = 300  # 5 mins
 BUILDING_ELEMENT_TREE_TIMEOUT_MS = 60 * 1000  # 1 minute
 BROWSER_CLOSE_TIMEOUT = 180  # 3 minute
+BROWSER_DOWNLOAD_MAX_WAIT_TIME = 120  # 2 minute
 BROWSER_DOWNLOAD_TIMEOUT = 600  # 10 minute
 DOWNLOAD_FILE_PREFIX = "downloads"
 SAVE_DOWNLOADED_FILES_TIMEOUT = 180
diff --git a/skyvern/webeye/actions/handler.py b/skyvern/webeye/actions/handler.py
index 9d8968d5..4d46c2b1 100644
--- a/skyvern/webeye/actions/handler.py
+++ b/skyvern/webeye/actions/handler.py
@@ -16,6 +16,7 @@ from pydantic import BaseModel
 from skyvern.config import settings
 from skyvern.constants import (
     AUTO_COMPLETION_POTENTIAL_VALUES_COUNT,
+    BROWSER_DOWNLOAD_MAX_WAIT_TIME,
     BROWSER_DOWNLOAD_TIMEOUT,
     DROPDOWN_MENU_MAX_DISTANCE,
     REPO_ROOT_DIR,
@@ -474,25 +475,23 @@ async def handle_click_to_download_file_action(
         )
         return [ActionFailure(e, download_triggered=False)]
 
-    # wait 5s to start downloading
-    LOG.info(
-        "Sleep for 5s to let download finish",
-        task_id=task.task_id,
-        step_id=step.step_id,
-        workflow_run_id=task.workflow_run_id,
-    )
-    await asyncio.sleep(5)
-    list_files_after = list_files_in_directory(download_dir)
-    LOG.info(
-        "Number of files in download directory after click",
-        num_downloaded_files_after=len(list_files_after),
-        download_dir=download_dir,
-        task_id=task.task_id,
-        step_id=step.step_id,
-        workflow_run_id=task.workflow_run_id,
-    )
+    try:
+        async with asyncio.timeout(BROWSER_DOWNLOAD_MAX_WAIT_TIME):
+            while True:
+                list_files_after = list_files_in_directory(download_dir)
+                LOG.info(
+                    "Number of files in download directory after click",
+                    num_downloaded_files_after=len(list_files_after),
+                    download_dir=download_dir,
+                    task_id=task.task_id,
+                    step_id=step.step_id,
+                    workflow_run_id=task.workflow_run_id,
+                )
+                if len(list_files_after) > len(list_files_before):
+                    break
+                await asyncio.sleep(1)
 
-    if len(list_files_after) <= len(list_files_before):
+    except asyncio.TimeoutError:
         LOG.warning(
             "No file to download after click",
             task_id=task.task_id,