@silverbulletmd/silverbullet 2.9.0 → 2.11.0

package/client/space_lua/eval.ts

@@ -7,7 +7,9 @@ import type {
   LuaTableField,
   NumericType,
 } from "./ast.ts";
+import type { Config } from "../config.ts";
 import { LuaAttribute } from "./ast.ts";
+import { budgetTick, LuaBudgetStopped } from "./budget.ts";
 import { evalPromiseValues } from "./util.ts";
 import {
   getMetatable,
@@ -91,6 +93,11 @@ import { getBlockGotoMeta } from "./labels.ts";
 
 const astNumberKindCache = new WeakMap<LuaExpression, NumericType>();
 
+function getClientConfig(): Config | undefined {
+  return (globalThis as typeof globalThis & { client?: { config: Config } })
+    .client?.config;
+}
+
 function astNumberKind(e: LuaExpression | undefined): NumericType | undefined {
   if (!e) return undefined;
 
@@ -258,7 +265,6 @@ export function luaOp(
 ): any {
   switch (op) {
     case "+": {
-      // Ultra-fast path: both plain numbers with no float type annotation (int + int)
       if (
         typeof left === "number" &&
         typeof right === "number" &&
@@ -308,11 +314,9 @@ export function luaOp(
       );
     }
     case "..": {
-      // Fast path: string .. string (most common in SilverBullet — key building, templates)
       if (typeof left === "string" && typeof right === "string") {
         return left + right;
       }
-      // Fast path: string .. number or number .. string
       if (typeof left === "string" && typeof right === "number") {
         return left + luaFormatNumber(right);
       }
@@ -352,7 +356,6 @@ export function luaOp(
       }
     }
     case "==": {
-      // Fast path for same-type primitives
       if (typeof left === typeof right && typeof left !== "object") {
         return left === right;
       }
@@ -370,11 +373,9 @@ export function luaOp(
       return !luaEqWithMetamethod(left, right, ctx, sf);
     }
     case "<": {
-      // Fast path: both plain numbers
       if (typeof left === "number" && typeof right === "number") {
         return left < right;
       }
-      // Fast path: both strings
       if (typeof left === "string" && typeof right === "string") {
         return left < right;
       }
@@ -404,7 +405,6 @@ export function luaOp(
     }
   }
 
-  // Remaining operators: //, %, bitwise
   const handler = operatorsMetaMethods[op];
   if (!handler) {
     throw new LuaRuntimeError(`Unknown operator ${op}`, sf.withCtx(ctx));
@@ -820,7 +820,6 @@ async function evalCrossJoinSources(
   sf: LuaStackFrame,
   ctx: ASTCtx,
 ): Promise<LuaTable[]> {
-  // Evaluate each source and normalize to arrays
   const arrays: { name: string; items: any[] }[] = [];
   for (const src of sources) {
     const val = await evalExpression(src.expression, env, sf);
@@ -834,7 +833,6 @@ async function evalCrossJoinSources(
     arrays.push({ name: src.name, items });
   }
 
-  // Cartesian product
   let product: Record<string, any>[] = [{}];
   for (const { name, items } of arrays) {
     const newProduct: Record<string, any>[] = [];
@@ -846,7 +844,6 @@ async function evalCrossJoinSources(
     product = newProduct;
   }
 
-  // Convert each combination to a `LuaTable` row
   return product.map((combo) => {
     const row = new LuaTable();
     for (const key in combo) {
@@ -894,7 +891,6 @@ export function evalExpression(
       case "Unary": {
         const u = asUnary(e);
 
-        // Fast path: negation of numeric literal
         if (u.operator === "-" && u.argument.type === "Number") {
           const num = u.argument;
           if (num.value === 0) {
@@ -914,7 +910,6 @@ export function evalExpression(
             const arg = singleResult(typed.value);
 
             return unaryWithMeta(arg, "__unm", u.ctx, sf, () => {
-              // Numeric-string coercion for unary minus
               if (typeof arg === "string") {
                 const n = coerceToNumber(arg);
                 if (n === null) {
@@ -1015,7 +1010,8 @@ export function evalExpression(
       case "Variable":
       case "FunctionCall":
       case "TableAccess":
-      case "PropertyAccess": {
+      case "PropertyAccess":
+      case "Parenthesized": {
         return evalPrefixExpression(e, env, sf);
       }
       case "TableConstructor": {
@@ -1108,7 +1104,6 @@ export function evalExpression(
         );
 
         if (fromSource.kind === "cross") {
-          // Materialize Cartesian product, then query
           return (async () => {
             const rows = await evalCrossJoinSources(
               fromSource.sources,
@@ -1118,13 +1113,11 @@ export function evalExpression(
             );
             const collection: any = toCollection(rows);
 
-            // Build up query object
             const query: LuaCollectionQuery = {
               objectVariable: undefined,
               distinct: true,
             };
 
-            // Map clauses to query parameters
             for (const clause of q.clauses) {
               switch (clause.type) {
                 case "Where": {
@@ -1178,12 +1171,11 @@ export function evalExpression(
             }
 
             return (collection as any)
-              .query(query, env, sf, globalThis.client?.config)
+              .query(query, env, sf, getClientConfig())
               .then(jsToLuaValue);
           })();
         }
 
-        // Single-source
         const { objectVariable, expression: objectExpression } = fromSource;
         return Promise.resolve(evalExpression(objectExpression, env, sf)).then(
           async (collection: LuaValue) => {
@@ -1199,33 +1191,27 @@ export function evalExpression(
               "query" in collection &&
               typeof (collection as any).query === "function"
             ) {
-              // Already queryable, use as-is
             } else if (collection instanceof LuaTable && collection.empty()) {
-              // Empty table → empty array
               collection = toCollection([]);
             } else if (collection instanceof LuaTable) {
               if (collection.length > 0) {
-                // Array-like table: extract array items, keep as LuaTables
                 const arr: any[] = [];
                 for (let i = 1; i <= collection.length; i++) {
                   arr.push(collection.rawGet(i));
                 }
                 collection = toCollection(arr);
               } else {
-                // Record-like table (no array part): treat as singleton
                 collection = toCollection([collection]);
               }
             } else {
               collection = toCollection(luaValueToJS(collection, sf));
             }
 
-            // Build up query object
             const query: LuaCollectionQuery = {
               objectVariable,
               distinct: true,
             };
 
-            // Map clauses to query parameters
             for (const clause of q.clauses) {
               switch (clause.type) {
                 case "Where": {
@@ -1278,9 +1264,8 @@ export function evalExpression(
               }
             }
 
-            // Always use the possibly-wrapped collection
             return (collection as any)
-              .query(query, env, sf, globalThis.client?.config)
+              .query(query, env, sf, getClientConfig())
               .then(jsToLuaValue);
           },
         );
@@ -1292,12 +1277,12 @@ export function evalExpression(
         );
     }
   } catch (err: any) {
-    // Repackage any non Lua-specific exceptions with some position information
-    if (!err.constructor.name.startsWith("Lua")) {
-      throw new LuaRuntimeError(err.message, sf.withCtx(e.ctx), err);
-    } else {
+    // instanceof, not constructor.name: the production build minifies class
+    // names, so a name-based test silently fails in the shipped client.
+    if (err instanceof LuaRuntimeError || err instanceof LuaBudgetStopped) {
       throw err;
     }
+    throw new LuaRuntimeError(err.message, sf.withCtx(e.ctx), err);
   }
 }
 
@@ -1318,10 +1303,9 @@ function evalPrefixExpression(
 
     case "Parenthesized": {
       const p = asParenthesized(e);
-      return evalExpression(p.expression, env, sf);
+      return rpThen(evalExpression(p.expression, env, sf), singleResult);
     }
 
-    // <<expr>>[<<expr>>]
     case "TableAccess": {
       const ta = asTableAccess(e);
       // Sync-first: evaluate object and key without allocating Promise when both are sync.
@@ -1341,7 +1325,6 @@ function evalPrefixExpression(
       );
     }
 
-    // <expr>.property
     case "PropertyAccess": {
       const pa = asPropertyAccess(e);
       // Sync-first: evaluate object; avoid Promise when object is sync.
@@ -1377,7 +1360,6 @@ function evalPrefixExpression(
         throw new LuaRuntimeError(nilMsg, sf.withCtx(fc.prefix.ctx));
       }
 
-      // Fast path: non-method call with sync prefix
       if (!fc.name && !isPromise(prefixValue)) {
         const argsVal = evalExpressions(fc.args, env, sf);
         if (!isPromise(argsVal)) {
@@ -1392,7 +1374,6 @@ function evalPrefixExpression(
         calleeVal: LuaValue,
         selfArgs: LuaValue[],
       ): LuaValue | Promise<LuaValue> => {
-        // Normal argument handling for hello:there(a, b, c) type calls
         if (fc.name) {
           const self = calleeVal;
           calleeVal = luaIndexValue(calleeVal, fc.name, sf);
@@ -1440,7 +1421,6 @@ function evalPrefixExpression(
   }
 }
 
-// Helper functions to reduce duplication
 function evalMetamethod(
   left: any,
   right: any,
@@ -1465,7 +1445,6 @@ function evalMetamethod(
   }
 }
 
-// Unary metamethod lookup and call
 function evalUnaryMetamethod(
   value: any,
   metaMethod: "__unm" | "__bnot",
@@ -1483,7 +1462,6 @@ function evalUnaryMetamethod(
   return luaCall(fn, [value], ctx, sf);
 }
 
-// Unary metamethod handling (with fallback)
 function unaryWithMeta(
   arg: any,
   meta: "__unm" | "__bnot",
@@ -1501,7 +1479,6 @@ function unaryWithMeta(
   return fallback();
 }
 
-// Logical short-circuit evaluation
 function evalLogical(
   op: "and" | "or",
   leftExpr: LuaExpression,
@@ -1770,12 +1747,10 @@ function luaRelWithMetamethod(
  * - throw error otherwise.
  */
 function luaLengthOp(val: any, ctx: ASTCtx, sf: LuaStackFrame): LuaValue {
-  // Strings: ignore `__len`
   if (typeof val === "string") {
     return val.length;
   }
 
-  // Tables: prefer metatable `__len` to raw length
   if (val instanceof LuaTable) {
     const mt = getMetatable(val, sf);
     if (mt) {
@@ -1787,7 +1762,6 @@ function luaLengthOp(val: any, ctx: ASTCtx, sf: LuaStackFrame): LuaValue {
     return val.length;
   }
 
-  // Other values: allow metatable `__len` first
   {
     const mt = getMetatable(val, sf);
     if (mt) {
@@ -1798,12 +1772,10 @@ function luaLengthOp(val: any, ctx: ASTCtx, sf: LuaStackFrame): LuaValue {
     }
   }
 
-  // JS arrays (interop): length if no `__len` override
   if (Array.isArray(val)) {
     return val.length;
   }
 
-  // Otherwise error with type
   const t = luaTypeOf(val) as LuaType;
   throw new LuaRuntimeError(
     `attempt to get length of a ${t} value`,
@@ -1819,7 +1791,6 @@ function evalExpressions(
   const len = es.length;
   if (len === 0) return [];
 
-  // Evaluate all arguments (sync-first); avoid .map() closure overhead
   const parts = new Array(len);
   for (let i = 0; i < len; i++) {
     parts[i] = evalExpression(es[i], env, sf);
@@ -1830,11 +1801,9 @@ function evalExpressions(
   const finalize = (argsResolved: any[]) => {
     const out: LuaValue[] = [];
     const lastIdx = argsResolved.length - 1;
-    // All but last expression produce a single value
     for (let i = 0; i < lastIdx; i++) {
       out.push(singleResult(argsResolved[i]));
     }
-    // Last expression preserves multiple results
     const last = argsResolved[lastIdx];
     if (last instanceof LuaMultiRes) {
       out.push(...last.flatten().values);
@@ -1866,6 +1835,8 @@ function runStatementsNoGoto(
       const result = evalStatement(stmts[i], execEnv, sf, returnOnReturn);
       if (isPromise(result)) {
         return (result as Promise<any>).then((res) => {
+          const budget = sf.threadState.budget;
+          if (budget !== undefined) budget.awaited = true;
           if (res !== undefined) {
             if (isGotoSignal(res)) {
               throw new LuaRuntimeError(
@@ -1894,6 +1865,35 @@ function runStatementsNoGoto(
   return processFrom(startIdx);
 }
 
+function closeThenRethrow(
+  sf: LuaStackFrame,
+  mark: number,
+  e: any,
+): never | Promise<never> {
+  const errObj: LuaValue =
+    e instanceof LuaRuntimeError ? e.message : (e?.message ?? String(e));
+  // Lua 5.4 lets an error from __close replace the one being unwound, but a
+  // user-initiated stop is not a Lua error and outranks it.
+  const preferOriginal = (closeErr: any): never => {
+    if (e instanceof LuaBudgetStopped) {
+      throw e;
+    }
+    throw closeErr;
+  };
+  let r: void | Promise<void>;
+  try {
+    r = luaCloseFromMark(sf, mark, errObj);
+  } catch (closeErr: any) {
+    return preferOriginal(closeErr);
+  }
+  if (isPromise(r)) {
+    return (r as Promise<void>).then(() => {
+      throw e;
+    }, preferOriginal);
+  }
+  throw e;
+}
+
 function withCloseBoundary(
   sf: LuaStackFrame,
   mark: number,
@@ -1914,17 +1914,7 @@ function withCloseBoundary(
     return isPromise(r) ? (r as Promise<void>).then(() => res) : res;
   };
 
-  const onRejected = (e: any) => {
-    const errObj: LuaValue =
-      e instanceof LuaRuntimeError ? e.message : (e?.message ?? String(e));
-    const r = luaCloseFromMark(sf, mark, errObj);
-    if (isPromise(r)) {
-      return (r as Promise<void>).then(() => {
-        throw e;
-      });
-    }
-    throw e;
-  };
+  const onRejected = (e: any) => closeThenRethrow(sf, mark, e);
 
   return p.then(onFulfilled, onRejected);
 }
@@ -1964,6 +1954,8 @@ function evalBlockNoClose(
         const r = evalStatement(stmts[i], execEnv, sf, returnOnReturn);
         if (isPromise(r)) {
           return (r as Promise<any>).then((res) => {
+            const budget = sf.threadState.budget;
+            if (budget !== undefined) budget.awaited = true;
             if (isGotoSignal(res)) return res;
             if (res !== undefined) return res;
             return runFrom(i + 1);
@@ -2004,11 +1996,13 @@ function evalBlockNoClose(
   const execEnv = b.needsEnv === true ? new LuaEnv(env) : env;
   const stmts = b.statements;
 
+  const budget = sf.threadState.budget;
   const runFrom = (i: number): EvalBlockResult => {
     for (; i < stmts.length; i++) {
       const r = evalStatement(stmts[i], execEnv, sf, returnOnReturn);
       if (isPromise(r)) {
         return (r as Promise<any>).then((res) => {
+          if (budget !== undefined) budget.awaited = true;
           const consumed = consumeGotoInBlock(res, meta!.labels);
           if (typeof consumed === "number") {
             return runFrom(consumed);
@@ -2021,6 +2015,13 @@ function evalBlockNoClose(
       }
       const consumed = consumeGotoInBlock(r, meta.labels);
       if (typeof consumed === "number") {
+        if (budget !== undefined && --budget.ticks <= 0) {
+          const y = budgetTick(budget);
+          if (y !== undefined) {
+            const resumeAt = consumed;
+            return rpThen(y, () => runFrom(resumeAt));
+          }
+        }
         i = consumed - 1;
         continue;
       }
@@ -2059,10 +2060,15 @@ export function evalStatement(
       );
 
       const apply = (values: LuaValue[], lvalues: { env: any; key: any }[]) => {
-        // Create the error-reporting frame once, not per-lvalue
         let errSf: LuaStackFrame | undefined;
         const ps: Promise<any>[] = [];
         for (let i = 0; i < lvalues.length; i++) {
+          if (
+            values[i] instanceof LuaFunction &&
+            typeof lvalues[i].key === "string"
+          ) {
+            values[i].info.name ??= lvalues[i].key;
+          }
           const r = luaSet(
             lvalues[i].env,
             lvalues[i].key,
@@ -2115,6 +2121,7 @@ export function evalStatement(
       }
 
       const bindOne = (name: any, v: LuaValue) => {
+        if (v instanceof LuaFunction) v.info.name ??= name.name;
         const isConst = name.attributes?.includes(LuaAttribute.Const) === true;
         const isClose = name.attributes?.includes(LuaAttribute.Close) === true;
 
@@ -2225,22 +2232,13 @@ export function evalStatement(
       try {
         out = evalBlockNoClose(b, env, sf, returnOnReturn);
       } catch (e: any) {
-        const errObj: LuaValue =
-          e instanceof LuaRuntimeError ? e.message : (e?.message ?? String(e));
-        const r = luaCloseFromMark(sf, mark, errObj);
-        if (isPromise(r)) {
-          return (r as Promise<void>).then(() => {
-            throw e;
-          });
-        }
-        throw e;
+        return closeThenRethrow(sf, mark, e);
       }
 
       return withCloseBoundary(sf, mark, out);
     }
     case "If": {
       const iff = asIf(s);
-      // Evaluate conditions in order; avoid awaiting when not necessary
       const conds = iff.conditions;
 
       const runFrom = (
@@ -2271,6 +2269,7 @@ export function evalStatement(
     }
     case "While": {
       const w = asWhile(s);
+      const budget = sf.threadState.budget;
 
       // Sync-first loop that re-enters sync mode after each async iteration
       const runSyncFirst = ():
@@ -2278,9 +2277,14 @@ export function evalStatement(
         | ControlSignal
         | Promise<undefined | ControlSignal> => {
         while (true) {
+          if (budget !== undefined && --budget.ticks <= 0) {
+            const y = budgetTick(budget);
+            if (y !== undefined) return rpThen(y, () => runSyncFirst());
+          }
           const c = evalExpression(w.condition, env, sf);
           if (isPromise(c)) {
             return (c as Promise<any>).then((cv) => {
+              if (budget !== undefined) budget.awaited = true;
               if (!luaTruthy(cv)) return;
               return rpThen(
                 evalStatement(w.block, env, sf, returnOnReturn),
@@ -2297,6 +2301,7 @@ export function evalStatement(
           const r = evalStatement(w.block, env, sf, returnOnReturn);
           if (isPromise(r)) {
             return (r as Promise<any>).then((res) => {
+              if (budget !== undefined) budget.awaited = true;
               if (res !== undefined) {
                 return isBreakSignal(res) ? undefined : res;
               }
@@ -2315,6 +2320,7 @@ export function evalStatement(
     }
     case "Repeat": {
       const rep = asRepeat(s);
+      const budget = sf.threadState.budget;
 
       // Sync-first loop that re-enters sync mode after each async iteration
       const runSyncFirst = ():
@@ -2322,9 +2328,14 @@ export function evalStatement(
         | ControlSignal
         | Promise<undefined | ControlSignal> => {
         while (true) {
+          if (budget !== undefined && --budget.ticks <= 0) {
+            const y = budgetTick(budget);
+            if (y !== undefined) return rpThen(y, () => runSyncFirst());
+          }
           const rr = evalStatement(rep.block, env, sf, returnOnReturn);
           if (isPromise(rr)) {
             return (rr as Promise<any>).then((res) => {
+              if (budget !== undefined) budget.awaited = true;
               if (res !== undefined) {
                 return isBreakSignal(res) ? undefined : res;
               }
@@ -2340,9 +2351,10 @@ export function evalStatement(
 
           const c = evalExpression(rep.condition, env, sf);
           if (isPromise(c)) {
-            return (c as Promise<any>).then((cv) =>
-              luaTruthy(cv) ? undefined : runSyncFirst(),
-            );
+            return (c as Promise<any>).then((cv) => {
+              if (budget !== undefined) budget.awaited = true;
+              return luaTruthy(cv) ? undefined : runSyncFirst();
+            });
           }
           if (luaTruthy(c)) break;
         }
@@ -2367,7 +2379,6 @@ export function evalStatement(
       let body = fn.body;
       let propNames = fn.name.propNames;
       if (fn.name.colonName) {
-        // function hello:there() -> function hello.there(self) transformation
         body = {
           ...fn.body,
           parameters: ["self", ...fn.body.parameters],
@@ -2386,13 +2397,13 @@ export function evalStatement(
       }
       (settable as any).set(
         propNames[propNames.length - 1],
-        new LuaFunction(body, env),
+        new LuaFunction(body, env, { name: propNames.join(".") }),
       );
       return;
     }
     case "LocalFunction": {
       const lf = asLocalFunction(s);
-      env.setLocal(lf.name, new LuaFunction(lf.body, env));
+      env.setLocal(lf.name, new LuaFunction(lf.body, env, { name: lf.name }));
       return;
     }
     case "Return": {
@@ -2494,10 +2505,22 @@ export function evalStatement(
         const shouldContinue =
           step > 0 ? (i: number) => i <= end : (i: number) => i >= end;
 
+        const budget = sf.threadState.budget;
+
         for (let i = startIndex; shouldContinue(i); i += step) {
+          if (budget !== undefined && --budget.ticks <= 0) {
+            const y = budgetTick(budget);
+            if (y !== undefined) {
+              const resumeAt = i;
+              return rpThen(y, () =>
+                runFromIndex(loopEnv, end, step, resumeAt, loopType),
+              );
+            }
+          }
           const r = executeIteration(loopEnv, i, loopType);
           if (isPromise(r)) {
             return (r as Promise<any>).then((res) => {
+              if (budget !== undefined) budget.awaited = true;
               if (res !== undefined) {
                 return isBreakSignal(res) ? undefined : res;
               }
@@ -2526,11 +2549,22 @@ export function evalStatement(
           step > 0 ? (i: number) => i <= end : (i: number) => i >= end;
 
         const loopEnv = new LuaEnv(env);
+        const budget = sf.threadState.budget;
 
         for (let i = start; shouldContinue(i); i += step) {
+          if (budget !== undefined && --budget.ticks <= 0) {
+            const y = budgetTick(budget);
+            if (y !== undefined) {
+              const resumeAt = i;
+              return rpThen(y, () =>
+                runSyncFirst(resumeAt, end, step, loopType),
+              );
+            }
+          }
           const r = executeIteration(loopEnv, i, loopType);
           if (isPromise(r)) {
             return (r as Promise<any>).then((res) => {
+              if (budget !== undefined) budget.awaited = true;
               if (res !== undefined) {
                 if (isBreakSignal(res)) {
                   return;
@@ -2565,6 +2599,8 @@ export function evalStatement(
         isPromise(endV) ? endV : Promise.resolve(endV),
         isPromise(stepV) ? stepV : Promise.resolve(stepV),
       ]).then(([start, end, step]) => {
+        const budget = sf.threadState.budget;
+        if (budget !== undefined) budget.awaited = true;
         return runSyncFirst(
           untagNumber(start) as number,
           untagNumber(end) as number,
@@ -2594,8 +2630,6 @@ export function evalStatement(
       const afterExprs = (resolved: any[]) => {
         const iteratorMultiRes = new LuaMultiRes(resolved).flatten();
         let iteratorValue: ILuaFunction | any = iteratorMultiRes.values[0];
-        // Handle the case where the iterator is a table and we need
-        // to call the `each` function.
         if (Array.isArray(iteratorValue) || iteratorValue instanceof LuaTable) {
           iteratorValue = (env.get("each") as ILuaFunction).call(
             sf,
@@ -2622,26 +2656,14 @@ export function evalStatement(
           luaMarkToBeClosed(sf, closing, fi.ctx);
         }
 
-        const errObjFrom = (e: any): LuaValue =>
-          e instanceof LuaRuntimeError ? e.message : (e?.message ?? String(e));
-
         const finish = (res: any) => {
           const r = luaCloseFromMark(sf, mark, null);
           return isPromise(r) ? (r as Promise<void>).then(() => res) : res;
         };
 
-        const finishErr = (e: any): Promise<never> | never => {
-          const errObj = errObjFrom(e);
-          const r = luaCloseFromMark(sf, mark, errObj);
-          if (isPromise(r)) {
-            return (r as Promise<void>).then(() => {
-              throw e;
-            });
-          }
-          throw e;
-        };
+        const finishErr = (e: any): Promise<never> | never =>
+          closeThenRethrow(sf, mark, e);
 
-        // Allocate the reusable env once before the loop
         const loopEnv = canReuseEnv ? new LuaEnv(env) : null;
 
         const makeIterEnv = (): LuaEnv => {
@@ -2651,9 +2673,14 @@ export function evalStatement(
           return new LuaEnv(env);
         };
 
+        const budget = sf.threadState.budget;
         // Sync-first loop that re-enters sync mode after each async iteration
         const runSyncFirst = (): any => {
           while (true) {
+            if (budget !== undefined && --budget.ticks <= 0) {
+              const y = budgetTick(budget);
+              if (y !== undefined) return rpThen(y, () => runSyncFirst());
+            }
             const iterCall = luaCall(
               iteratorValue,
               [state, control],
@@ -2662,6 +2689,7 @@ export function evalStatement(
             );
 
             const afterIterCall = (itv: any): any => {
+              if (budget !== undefined) budget.awaited = true;
               const iterResult = new LuaMultiRes(itv).flatten();
               const nextControl = iterResult.values[0];
               if (nextControl === null || nextControl === undefined) {
@@ -2704,6 +2732,7 @@ export function evalStatement(
             if (isPromise(r)) {
               return (r as Promise<any>)
                 .then((res) => {
+                  if (budget !== undefined) budget.awaited = true;
                   if (res !== undefined) {
                     if (isBreakSignal(res)) {
                       return finish(undefined);
@@ -2731,7 +2760,11 @@ export function evalStatement(
       };
 
       if (isPromise(exprVals)) {
-        return (exprVals as Promise<any[]>).then(afterExprs);
+        return (exprVals as Promise<any[]>).then((v) => {
+          const budget = sf.threadState.budget;
+          if (budget !== undefined) budget.awaited = true;
+          return afterExprs(v);
+        });
       }
       return afterExprs(exprVals as any[]);
     }