complete_on_downloads for task block (#403)

skyvern/webeye/actions/handler.py

@@ -13,7 +13,11 @@ from skyvern.constants import REPO_ROOT_DIR
 from skyvern.exceptions import ImaginaryFileUrl, MissingElement, MissingFileUrl, MultipleElementsFound
 from skyvern.forge import app
 from skyvern.forge.prompts import prompt_engine
-from skyvern.forge.sdk.api.files import download_file
+from skyvern.forge.sdk.api.files import (
+    download_file,
+    get_number_of_files_in_directory,
+    get_path_for_workflow_download_directory,
+)
 from skyvern.forge.sdk.models import Step
 from skyvern.forge.sdk.schemas.tasks import Task
 from skyvern.forge.sdk.services.bitwarden import BitwardenConstants
@@ -162,17 +166,42 @@ async def handle_click_action(
     task: Task,
     step: Step,
 ) -> list[ActionResult]:
+    num_downloaded_files_before = 0
+    download_dir = None
+    if task.workflow_run_id:
+        download_dir = get_path_for_workflow_download_directory(task.workflow_run_id)
+        num_downloaded_files_before = get_number_of_files_in_directory(download_dir)
+        LOG.info(
+            "Number of files in download directory before click",
+            num_downloaded_files_before=num_downloaded_files_before,
+            download_dir=download_dir,
+        )
     xpath = await validate_actions_in_dom(action, page, scraped_page)
     await asyncio.sleep(0.3)
     if action.download:
-        return await handle_click_to_download_file_action(action, page, scraped_page)
-    return await chain_click(
-        task,
-        page,
-        action,
-        xpath,
-        timeout=SettingsManager.get_settings().BROWSER_ACTION_TIMEOUT_MS,
-    )
+        results = await handle_click_to_download_file_action(action, page, scraped_page)
+    else:
+        results = await chain_click(
+            task,
+            page,
+            action,
+            xpath,
+            timeout=SettingsManager.get_settings().BROWSER_ACTION_TIMEOUT_MS,
+        )
+
+    if results and task.workflow_run_id and download_dir:
+        LOG.info("Sleeping for 5 seconds to let the download finish")
+        await asyncio.sleep(5)
+        num_downloaded_files_after = get_number_of_files_in_directory(download_dir)
+        LOG.info(
+            "Number of files in download directory after click",
+            num_downloaded_files_after=num_downloaded_files_after,
+            download_dir=download_dir,
+        )
+        if num_downloaded_files_after > num_downloaded_files_before:
+            results[-1].download_triggered = True
+
+    return results
 
 
 async def handle_click_to_download_file_action(
@@ -677,16 +706,6 @@ async def chain_click(
     page.on("filechooser", fc_func)
     LOG.info("Registered file chooser listener", action=action, path=file)
 
-    # If a download is triggered due to the click, we need to let LLM know in action_results
-    download_triggered = False
-
-    def download_func(download: Any) -> None:
-        nonlocal download_triggered
-        download_triggered = True
-
-    page.on("download", download_func)
-    LOG.info("Registered download listener", action=action)
-
     """
     Clicks on an element identified by the xpath and its parent if failed.
     :param xpath: xpath of the element to click
@@ -698,7 +717,6 @@ async def chain_click(
         return [
             ActionSuccess(
                 javascript_triggered=javascript_triggered,
-                download_triggered=download_triggered,
             )
         ]
     except Exception as e:
@@ -706,7 +724,6 @@ async def chain_click(
             ActionFailure(
                 e,
                 javascript_triggered=javascript_triggered,
-                download_triggered=download_triggered,
             )
         ]
         if await is_input_element(page.locator(xpath)):
@@ -716,7 +733,6 @@ async def chain_click(
                 xpath=xpath,
             )
             sibling_action_result = await click_sibling_of_input(page.locator(xpath), timeout=timeout)
-            sibling_action_result.download_triggered = download_triggered
             action_results.append(sibling_action_result)
             if type(sibling_action_result) == ActionSuccess:
                 return action_results
@@ -736,7 +752,6 @@ async def chain_click(
                 ActionSuccess(
                     javascript_triggered=javascript_triggered,
                     interacted_with_parent=True,
-                    download_triggered=download_triggered,
                 )
             )
         except Exception as pe:
@@ -765,7 +780,6 @@ async def chain_click(
         if file:
             await asyncio.sleep(10)
         page.remove_listener("filechooser", fc_func)
-        page.remove_listener("download", download_func)
 
 
 def get_anchor_to_click(scraped_page: ScrapedPage, element_id: int) -> str | None: