Fix infinite loop detection false trigger on await in loops (#4295)

client/modules/Preview/jsPreprocess.js

@@ -39,14 +39,27 @@ function makeVarDecl(varName) {
         type: 'VariableDeclarator',
         id: { type: 'Identifier', name: varName },
         init: {
-          type: 'CallExpression',
-          callee: {
-            type: 'MemberExpression',
-            object: { type: 'Identifier', name: 'Date' },
-            property: { type: 'Identifier', name: 'now' },
-            computed: false
-          },
-          arguments: []
+          type: 'ObjectExpression',
+          properties: [
+            {
+              type: 'Property',
+              key: { type: 'Identifier', name: 't' },
+              value: {
+                type: 'CallExpression',
+                callee: {
+                  type: 'MemberExpression',
+                  object: { type: 'Identifier', name: 'Date' },
+                  property: { type: 'Identifier', name: 'now' },
+                  computed: false
+                },
+                arguments: []
+              },
+              kind: 'init',
+              method: false,
+              shorthand: false,
+              computed: false
+            }
+          ]
         }
       }
     ]
@@ -72,7 +85,12 @@ function makeCheckStatement(varName, line) {
           },
           arguments: []
         },
-        right: { type: 'Identifier', name: varName }
+        right: {
+          type: 'MemberExpression',
+          object: { type: 'Identifier', name: varName },
+          property: { type: 'Identifier', name: 't' },
+          computed: false
+        }
       },
       right: {
         type: 'Literal',
@@ -160,6 +178,8 @@ function collectLoopsToProtect(ast, shaderNames) {
 
   walk.ancestor(ast, {
     ForStatement: visitNode,
+    ForInStatement: visitNode,
+    ForOfStatement: visitNode,
     WhileStatement: visitNode,
     DoWhileStatement: visitNode
   });
@@ -179,6 +199,41 @@ function injectProtection(loops) {
       loop.body = { type: 'BlockStatement', body: [check, loop.body] };
     }
 
+    walk.ancestor(loop.body, {
+      AwaitExpression(node, ancestors) {
+        const isInsideNestedFunction = ancestors.some(
+          (ancestor) =>
+            ancestor !== loop.body &&
+            (ancestor.type === 'FunctionDeclaration' ||
+              ancestor.type === 'FunctionExpression' ||
+              ancestor.type === 'ArrowFunctionExpression')
+        );
+
+        if (!isInsideNestedFunction) {
+          const originalAwait = {
+            type: 'AwaitExpression',
+            argument: node.argument
+          };
+          node.type = 'CallExpression';
+          node.callee = {
+            type: 'MemberExpression',
+            object: {
+              type: 'MemberExpression',
+              object: { type: 'Identifier', name: 'window' },
+              property: { type: 'Identifier', name: 'loopProtect' },
+              computed: false
+            },
+            property: { type: 'Identifier', name: 'reset' },
+            computed: false
+          };
+          node.arguments = [
+            { type: 'Identifier', name: varName },
+            originalAwait
+          ];
+        }
+      }
+    });
+
     if (parentBlock) {
       const varDecl = makeVarDecl(varName);
       const nodeIdx = parentBlock.body.indexOf(loop);

client/modules/Preview/jsPreprocess.unit.test.js

@@ -148,4 +148,51 @@ describe('jsPreprocess', () => {
       expect(result).toContain('window.loopProtect.hit');
     });
   });
+
+  describe('async/await loop protection', () => {
+    it('wraps await expressions in for loop with window.loopProtect.reset', () => {
+      const code = `
+        async function setup() {
+          for (let i = 0; i < fns.length; i++) {
+            imgs.push(await loadImage(fns[i]));
+          }
+        }
+      `;
+      const result = jsPreprocess(code, '');
+      expect(result).toContain('window.loopProtect.hit');
+      expect(result).toContain(
+        'window.loopProtect.reset(_LP0, await loadImage(fns[i]))'
+      );
+    });
+
+    it('does not wrap await expressions outside loops', () => {
+      const code = `
+        async function setup() {
+          const img = await loadImage('a.jpg');
+          for (let i = 0; i < 10; i++) {}
+        }
+      `;
+      const result = jsPreprocess(code, '');
+      expect(result).toContain("await loadImage('a.jpg')");
+      expect(result).not.toContain(
+        'window.loopProtect.reset(_LP0, await loadImage'
+      );
+    });
+
+    it('handles nested loops with await expressions', () => {
+      const code = `
+        async function run() {
+          for (let i = 0; i < 2; i++) {
+            for (let j = 0; j < 2; j++) {
+              await loadData(i, j);
+            }
+          }
+        }
+      `;
+      const result = jsPreprocess(code, '');
+      expect(result).toContain(
+        'window.loopProtect.reset(_LP0, window.loopProtect.reset(_LP1, await loadData(i, j)))'
+      );
+    });
+  });
 });

client/utils/previewEntry.js

@@ -51,6 +51,12 @@ window.loopProtect = {
     }
 
     return true;
+  },
+  reset: function resetLoopTimer(loopObj, val) {
+    if (loopObj) {
+      loopObj.t = Date.now();
+    }
+    return val;
   }
 };