diff --git a/src/ast/parsePrefix.zig b/src/ast/parsePrefix.zig
index 21353daae0..69f19068ca 100644
--- a/src/ast/parsePrefix.zig
+++ b/src/ast/parsePrefix.zig
@@ -136,9 +136,14 @@ pub fn ParsePrefix(
                         .allow_ident => {
                             p.lexer.prev_token_was_await_keyword = true;
                             p.lexer.await_keyword_loc = name_range.loc;
-                            p.lexer.fn_or_arrow_start_loc = p.fn_or_arrow_data_parse.needs_async_loc;
-                            // Check if we're at module scope by comparing current scope to module scope
-                            p.lexer.is_at_module_scope = p.current_scope == p.module_scope;
+                            // Use fn_or_arrow_start_loc as a signal:
+                            // - If at module scope, set it to await_keyword_loc (sentinel)
+                            // - Otherwise, set it to needs_async_loc (may be empty for arrow functions)
+                            if (p.current_scope == p.module_scope) {
+                                p.lexer.fn_or_arrow_start_loc = name_range.loc;
+                            } else {
+                                p.lexer.fn_or_arrow_start_loc = p.fn_or_arrow_data_parse.needs_async_loc;
+                            }
                         },
                     }
                 },
diff --git a/src/js_lexer.zig b/src/js_lexer.zig
index 008a084952..2c16748e3a 100644
--- a/src/js_lexer.zig
+++ b/src/js_lexer.zig
@@ -138,7 +138,6 @@ fn NewLexer_(
         prev_token_was_await_keyword: bool = false,
         await_keyword_loc: logger.Loc = logger.Loc.Empty,
         fn_or_arrow_start_loc: logger.Loc = logger.Loc.Empty,
-        is_at_module_scope: bool = true,
         regex_flags_start: ?u16 = null,
         allocator: std.mem.Allocator,
         string_literal_raw_content: string = "",
@@ -1821,7 +1820,13 @@ fn NewLexer_(
 
         pub fn expectedString(self: *LexerType, text: string) !void {
             if (self.prev_token_was_await_keyword) {
-                if (self.is_at_module_scope) {
+                // Use fn_or_arrow_start_loc as a signal:
+                // - If it equals await_keyword_loc, we're at module scope (sentinel value)
+                // - If it's not empty and different, we're in a function with a location for "async"
+                // - If it's empty, we're in a function without a clear location (e.g., arrow function)
+                const is_at_module_scope = self.fn_or_arrow_start_loc.start == self.await_keyword_loc.start;
+
+                if (is_at_module_scope) {
                     // Top-level await - provide better error message
                     try self.addRangeErrorWithNotes(
                         self.range(),
diff --git a/test/bundler/bundler_top_level_await_error.test.ts b/test/bundler/bundler_top_level_await_error.test.ts
index 02233c099b..b3484bd712 100644
--- a/test/bundler/bundler_top_level_await_error.test.ts
+++ b/test/bundler/bundler_top_level_await_error.test.ts
@@ -54,16 +54,14 @@ describe("bundler", () => {
   itBundled("top_level_await_error/AwaitInAsyncFunctionShouldStillWork", {
     files: {
       "/entry.ts": /* ts */ `
-        async function main() {
+        (async function main() {
           async function sum(a: number, b: number) {
             return a + b;
           }
 
           const result = await sum(5, 5);
           console.log(result);
-        }
-
-        main().catch(console.error);
+        })().catch(console.error);
       `,
     },
     format: "cjs",