From 1bb76759fb52994d2b3eb4baf3908e53397186aa Mon Sep 17 00:00:00 2001
From: mondaychen <monday.chen@gmail.com>
Date: Mon, 8 Apr 2024 16:14:21 -0400
Subject: [PATCH] refactor: better way to ensure page has been loaded

---
 src/helpers/rpc/domActions.ts |  2 +-
 src/helpers/rpc/utils.ts      | 10 ++++++++++
 src/state/currentTask.ts      | 26 ++++++++++++++++++--------
 3 files changed, 29 insertions(+), 9 deletions(-)
 create mode 100644 src/helpers/rpc/utils.ts

diff --git a/src/helpers/rpc/domActions.ts b/src/helpers/rpc/domActions.ts
index 1daa048..be2b7ba 100644
--- a/src/helpers/rpc/domActions.ts
+++ b/src/helpers/rpc/domActions.ts
@@ -4,7 +4,7 @@ import { scrollScriptString } from "./runtimeFunctionStrings";
 import { sleep, waitFor, waitTillStable } from "../utils";
 
 const DEFAULT_INTERVAL = 500;
-const DEFAULT_TIMEOUT = 0;
+const DEFAULT_TIMEOUT = 10000; // 10 seconds
 
 export class DomActions {
   static delayBetweenClicks = 500; // Set this value to control the delay between clicks
diff --git a/src/helpers/rpc/utils.ts b/src/helpers/rpc/utils.ts
new file mode 100644
index 0000000..e32d648
--- /dev/null
+++ b/src/helpers/rpc/utils.ts
@@ -0,0 +1,10 @@
+import { DomActions } from "./domActions";
+
+export async function waitTillHTMLRendered(
+  tabId: number,
+  interval = undefined,
+  timeout = undefined,
+) {
+  const domActions = new DomActions(tabId);
+  return await domActions.waitTillHTMLRendered(interval, timeout);
+}
diff --git a/src/state/currentTask.ts b/src/state/currentTask.ts
index 95e68a9..a196534 100644
--- a/src/state/currentTask.ts
+++ b/src/state/currentTask.ts
@@ -16,11 +16,12 @@ import {
 } from "../helpers/vision-agent/parseResponse";
 import { callRPCWithTab } from "../helpers/rpc/pageRPC";
 import { getSimplifiedDom } from "../helpers/simplifyDom";
-import { sleep, truthyFilter } from "../helpers/utils";
+import { sleep, truthyFilter, waitFor } from "../helpers/utils";
 import {
   operateTool,
   operateToolWithSimpliedDom,
 } from "../helpers/rpc/performAction";
+import { waitTillHTMLRendered } from "../helpers/rpc/utils";
 import { findActiveTab } from "../helpers/browserUtils";
 import { MyStateCreator } from "./store";
 import buildAnnotatedScreenshots from "../helpers/buildAnnotatedScreenshots";
@@ -103,12 +104,24 @@ export const createCurrentTaskSlice: MyStateCreator<CurrentTaskSlice> = (
         while (true) {
           if (wasStopped()) break;
 
-          // get latest tab info, since button clicking might have changed it
-          const activeTab = await findActiveTab();
+          // always get latest tab info, since actions such as button clicking might have changed it
+          let activeTab = await findActiveTab();
           const tabId = activeTab?.id || -1;
           if (!activeTab || !tabId) {
             throw new Error("No active tab found");
           }
+          if (activeTab.status === "loading") {
+            // wait for tab to be loaded
+            await waitFor(
+              async () => {
+                // findActiveTab give a new reference to activeTab every time
+                activeTab = await findActiveTab();
+                return activeTab?.status === "complete";
+              },
+              200, // check every 200ms
+              100, // wait for up to 20 seconds (100*200ms)
+            );
+          }
 
           const isVisionModel = hasVisionSupport(get().settings.selectedModel);
 
@@ -168,14 +181,13 @@ export const createCurrentTaskSlice: MyStateCreator<CurrentTaskSlice> = (
             if (shouldContinue) {
               // if navigation was successful, continue the task on the new page
               setActionStatus("waiting");
-              // sleep 2 seconds. This is pretty arbitrary; we should figure out a better way to determine when the page has settled.
-              await sleep(2000);
               continue;
             } else {
               break;
             }
           }
           await attachDebugger(tabId);
+          await waitTillHTMLRendered(tabId);
 
           set((state) => {
             state.currentTask.tabId = tabId;
@@ -243,9 +255,7 @@ export const createCurrentTaskSlice: MyStateCreator<CurrentTaskSlice> = (
           }
 
           setActionStatus("waiting");
-          // sleep 2 seconds. This is pretty arbitrary; we should figure out a better way to determine when the page has settled.
-          await sleep(2000);
-        }
+        } // end of while loop
         set((state) => {
           state.currentTask.status = "success";
         });