diff --git a/src/ast.zig b/src/ast.zig
index 9e940f934..508683b69 100644
--- a/src/ast.zig
+++ b/src/ast.zig
@@ -1217,6 +1217,14 @@ pub fn isContainer(tree: Ast, node: Ast.Node.Index) bool {
     };
 }
 
+pub fn rootDecls(tree: Ast) []const Node.Index {
+    const nodes_data = tree.nodes.items(.data);
+    switch (tree.mode) {
+        .zig => return tree.extra_data[nodes_data[0].lhs..nodes_data[0].rhs],
+        .zon => return (&nodes_data[0].lhs)[0..1],
+    }
+}
+
 pub fn isBuiltinCall(tree: Ast, node: Ast.Node.Index) bool {
     return switch (tree.nodes.items(.tag)[node]) {
         .builtin_call,
diff --git a/src/features/code_actions.zig b/src/features/code_actions.zig
index c379127c0..bd08c8030 100644
--- a/src/features/code_actions.zig
+++ b/src/features/code_actions.zig
@@ -595,7 +595,7 @@ pub fn getImportsDecls(builder: *Builder, allocator: std.mem.Allocator) error{Ou
     const node_data = tree.nodes.items(.data);
     const node_tokens = tree.nodes.items(.main_token);
 
-    const root_decls = tree.rootDecls();
+    const root_decls = ast.rootDecls(tree);
 
     var skip_set = try std.DynamicBitSetUnmanaged.initEmpty(allocator, root_decls.len);
     defer skip_set.deinit(allocator);
diff --git a/src/features/diagnostics.zig b/src/features/diagnostics.zig
index 3028efb43..48d7ac037 100644
--- a/src/features/diagnostics.zig
+++ b/src/features/diagnostics.zig
@@ -79,7 +79,7 @@ pub fn generateDiagnostics(server: *Server, arena: std.mem.Allocator, handle: *D
 
         // TODO: style warnings for types, values and declarations below root scope
         if (tree.errors.len == 0) {
-            for (tree.rootDecls()) |decl_idx| {
+            for (ast.rootDecls(tree)) |decl_idx| {
                 const decl = tree.nodes.items(.tag)[decl_idx];
                 switch (decl) {
                     .fn_proto,
@@ -149,7 +149,7 @@ pub fn generateDiagnostics(server: *Server, arena: std.mem.Allocator, handle: *D
     if (server.config.highlight_global_var_declarations) {
         const main_tokens = tree.nodes.items(.main_token);
         const tags = tree.tokens.items(.tag);
-        for (tree.rootDecls()) |decl| {
+        for (ast.rootDecls(tree)) |decl| {
             const decl_tag = tree.nodes.items(.tag)[decl];
             const decl_main_token = tree.nodes.items(.main_token)[decl];
 
diff --git a/src/features/semantic_tokens.zig b/src/features/semantic_tokens.zig
index 24a442be9..9b58dccfb 100644
--- a/src/features/semantic_tokens.zig
+++ b/src/features/semantic_tokens.zig
@@ -1040,9 +1040,9 @@ pub fn writeSemanticTokens(
         .limited = limited,
     };
 
-    var nodes = if (loc) |l| try ast.nodesAtLoc(arena, handle.tree, l) else handle.tree.rootDecls();
+    var nodes = if (loc) |l| try ast.nodesAtLoc(arena, handle.tree, l) else ast.rootDecls(handle.tree);
     if (nodes.len == 1 and nodes[0] == 0) {
-        nodes = handle.tree.rootDecls();
+        nodes = ast.rootDecls(handle.tree);
     }
 
     // reverse the ast from the root declarations
diff --git a/tests/context.zig b/tests/context.zig
index f261dee04..d2ad8ec5f 100644
--- a/tests/context.zig
+++ b/tests/context.zig
@@ -69,15 +69,19 @@ pub const Context = struct {
     }
 
     // helper
-    pub fn addDocument(self: *Context, source: []const u8) ![]const u8 {
+    pub fn addDocument(self: *Context, options: struct {
+        uri: ?[]const u8 = null,
+        source: []const u8,
+        mode: std.zig.Ast.Mode = .zig,
+    }) ![]const u8 {
         const fmt = switch (builtin.os.tag) {
-            .windows => "file:///C:\\nonexistent\\test-{d}.zig",
-            else => "file:///nonexistent/test-{d}.zig",
+            .windows => "file:///C:\\nonexistent\\test-{d}.{s}",
+            else => "file:///nonexistent/test-{d}.{s}",
         };
-        const uri = try std.fmt.allocPrint(
+        const uri = options.uri orelse try std.fmt.allocPrint(
             self.arena.allocator(),
             fmt,
-            .{self.file_id},
+            .{ self.file_id, @tagName(options.mode) },
         );
 
         const params = types.DidOpenTextDocumentParams{
@@ -85,7 +89,7 @@ pub const Context = struct {
                 .uri = uri,
                 .languageId = "zig",
                 .version = 420,
-                .text = source,
+                .text = options.source,
             },
         };
 
diff --git a/tests/lsp_features/code_actions.zig b/tests/lsp_features/code_actions.zig
index 27716be50..5e2ec0248 100644
--- a/tests/lsp_features/code_actions.zig
+++ b/tests/lsp_features/code_actions.zig
@@ -596,7 +596,7 @@ fn testAutofixOptions(before: []const u8, after: []const u8, want_zir: bool) !vo
     ctx.server.config.enable_autofix = true;
     ctx.server.config.prefer_ast_check_as_child_process = !want_zir;
 
-    const uri = try ctx.addDocument(before);
+    const uri = try ctx.addDocument(.{ .source = before });
     const handle = ctx.server.document_store.getHandle(uri).?;
 
     var diagnostics: std.ArrayListUnmanaged(types.Diagnostic) = .{};
@@ -643,7 +643,7 @@ fn testDiagnostic(before: []const u8, after: []const u8) !void {
     defer ctx.deinit();
     ctx.server.config.enable_autofix = true;
 
-    const uri = try ctx.addDocument(before);
+    const uri = try ctx.addDocument(.{ .source = before });
     const handle = ctx.server.document_store.getHandle(uri).?;
 
     const params = types.CodeActionParams{
diff --git a/tests/lsp_features/completion.zig b/tests/lsp_features/completion.zig
index 6d8a03a60..5787d9240 100644
--- a/tests/lsp_features/completion.zig
+++ b/tests/lsp_features/completion.zig
@@ -3675,7 +3675,7 @@ fn testCompletionWithOptions(
     ctx.server.config.enable_snippets = options.enable_snippets;
     ctx.server.config.completion_label_details = options.completion_label_details;
 
-    const test_uri = try ctx.addDocument(text);
+    const test_uri = try ctx.addDocument(.{ .source = text });
 
     const params = types.CompletionParams{
         .textDocument = .{ .uri = test_uri },
@@ -3913,7 +3913,7 @@ fn testCompletionTextEdit(
     ctx.server.config.enable_argument_placeholders = options.enable_argument_placeholders;
     ctx.server.config.enable_snippets = options.enable_snippets;
 
-    const test_uri = try ctx.addDocument(text);
+    const test_uri = try ctx.addDocument(.{ .source = text });
     const handle = ctx.server.document_store.getHandle(test_uri).?;
 
     const cursor_position = offsets.indexToPosition(options.source, cursor_idx, ctx.server.offset_encoding);
diff --git a/tests/lsp_features/definition.zig b/tests/lsp_features/definition.zig
index a52615c1b..3f720dff6 100644
--- a/tests/lsp_features/definition.zig
+++ b/tests/lsp_features/definition.zig
@@ -233,7 +233,7 @@ fn testDefinition(source: []const u8) !void {
     var ctx = try Context.init();
     defer ctx.deinit();
 
-    const test_uri = try ctx.addDocument(phr.new_source);
+    const test_uri = try ctx.addDocument(.{ .source = phr.new_source });
 
     var error_builder = ErrorBuilder.init(allocator);
     defer error_builder.deinit();
diff --git a/tests/lsp_features/document_symbol.zig b/tests/lsp_features/document_symbol.zig
index 2c5678c6d..642da9dbe 100644
--- a/tests/lsp_features/document_symbol.zig
+++ b/tests/lsp_features/document_symbol.zig
@@ -97,7 +97,7 @@ fn testDocumentSymbol(source: []const u8, want: []const u8) !void {
     var ctx = try Context.init();
     defer ctx.deinit();
 
-    const test_uri = try ctx.addDocument(source);
+    const test_uri = try ctx.addDocument(.{ .source = source });
 
     const params = types.DocumentSymbolParams{
         .textDocument = .{ .uri = test_uri },
diff --git a/tests/lsp_features/folding_range.zig b/tests/lsp_features/folding_range.zig
index e24e7f0c1..b460a6220 100644
--- a/tests/lsp_features/folding_range.zig
+++ b/tests/lsp_features/folding_range.zig
@@ -302,7 +302,7 @@ fn testFoldingRange(source: []const u8, expect: []const types.FoldingRange) !voi
     var ctx = try Context.init();
     defer ctx.deinit();
 
-    const test_uri = try ctx.addDocument(source);
+    const test_uri = try ctx.addDocument(.{ .source = source });
 
     const params = types.FoldingRangeParams{ .textDocument = .{ .uri = test_uri } };
 
diff --git a/tests/lsp_features/hover.zig b/tests/lsp_features/hover.zig
index 24b9d502a..dbb84007c 100644
--- a/tests/lsp_features/hover.zig
+++ b/tests/lsp_features/hover.zig
@@ -1308,13 +1308,13 @@ fn testHoverWithOptions(
 
     ctx.server.client_capabilities.hover_supports_md = options.markup_kind == .markdown;
 
-    const test_uri = "file:///test.zig";
-    try ctx.server.sendNotificationSync(ctx.arena.allocator(), "textDocument/didOpen", .{
-        .textDocument = .{ .uri = test_uri, .languageId = "zig", .version = 420, .text = text },
+    const uri = try ctx.addDocument(.{
+        .uri = "file:///test.zig",
+        .source = text,
     });
 
     const params = types.HoverParams{
-        .textDocument = .{ .uri = test_uri },
+        .textDocument = .{ .uri = uri },
         .position = offsets.indexToPosition(text, cursor_idx, ctx.server.offset_encoding),
     };
 
diff --git a/tests/lsp_features/inlay_hints.zig b/tests/lsp_features/inlay_hints.zig
index 423246723..db4f622f4 100644
--- a/tests/lsp_features/inlay_hints.zig
+++ b/tests/lsp_features/inlay_hints.zig
@@ -509,7 +509,7 @@ fn testInlayHints(source: []const u8, options: Options) !void {
     ctx.server.config.inlay_hints_hide_redundant_param_names = options.hide_redundant_param_names;
     ctx.server.config.inlay_hints_hide_redundant_param_names_last_token = options.hide_redundant_param_names_last_token;
 
-    const test_uri = try ctx.addDocument(phr.new_source);
+    const test_uri = try ctx.addDocument(.{ .source = phr.new_source });
 
     const range = types.Range{
         .start = types.Position{ .line = 0, .character = 0 },
diff --git a/tests/lsp_features/references.zig b/tests/lsp_features/references.zig
index 30ad6d65f..017ad8be0 100644
--- a/tests/lsp_features/references.zig
+++ b/tests/lsp_features/references.zig
@@ -245,7 +245,7 @@ fn testMFReferences(sources: []const []const u8) !void {
         var phr = try helper.collectReplacePlaceholders(allocator, source, placeholder_name);
         defer phr.deinit(allocator);
 
-        const uri = try ctx.addDocument(phr.new_source);
+        const uri = try ctx.addDocument(.{ .source = phr.new_source });
         files.putAssumeCapacityNoClobber(uri, .{ .source = source, .new_source = phr.new_source });
         phr.new_source = ""; // `files` takes ownership of `new_source` from `phr`
 
diff --git a/tests/lsp_features/selection_range.zig b/tests/lsp_features/selection_range.zig
index a7b2480d7..069f05f4e 100644
--- a/tests/lsp_features/selection_range.zig
+++ b/tests/lsp_features/selection_range.zig
@@ -36,7 +36,7 @@ fn testSelectionRange(source: []const u8, want: []const []const u8) !void {
     var ctx = try Context.init();
     defer ctx.deinit();
 
-    const test_uri = try ctx.addDocument(phr.new_source);
+    const test_uri = try ctx.addDocument(.{ .source = phr.new_source });
 
     const position = offsets.locToRange(phr.new_source, phr.locations.items(.new)[0], .@"utf-16").start;
 
diff --git a/tests/lsp_features/semantic_tokens.zig b/tests/lsp_features/semantic_tokens.zig
index c41175e9b..549f9edbb 100644
--- a/tests/lsp_features/semantic_tokens.zig
+++ b/tests/lsp_features/semantic_tokens.zig
@@ -1727,6 +1727,25 @@ test "deprecated" {
     });
 }
 
+test "zon file" {
+    try testSemanticTokensOptions(
+        \\.{
+        \\    .foo = "bar",
+        \\    .baz = true,
+        \\}
+    , &.{
+        .{ ".", .property, .{} },
+        .{ "foo", .property, .{} },
+        .{ "=", .operator, .{} },
+        .{ "\"bar\"", .string, .{} },
+
+        .{ ".", .property, .{} },
+        .{ "baz", .property, .{} },
+        .{ "=", .operator, .{} },
+        .{ "true", .keywordLiteral, .{} },
+    }, .{ .mode = .zon });
+}
+
 test "recursive usingnamespace" {
     // this test is supposed to check against infinite recursion when resolving usingnamespace
     try testSemanticTokens(
@@ -1849,10 +1868,23 @@ const TokenIterator = struct {
 };
 
 fn testSemanticTokens(source: [:0]const u8, expected_tokens: []const TokenData) !void {
+    try testSemanticTokensOptions(source, expected_tokens, .{});
+}
+
+fn testSemanticTokensOptions(
+    source: [:0]const u8,
+    expected_tokens: []const TokenData,
+    options: struct {
+        mode: std.zig.Ast.Mode = .zig,
+    },
+) !void {
     var ctx = try Context.init();
     defer ctx.deinit();
 
-    const uri = try ctx.addDocument(source);
+    const uri = try ctx.addDocument(.{
+        .source = source,
+        .mode = options.mode,
+    });
 
     const params = types.SemanticTokensParams{
         .textDocument = .{ .uri = uri },
diff --git a/tests/lsp_features/signature_help.zig b/tests/lsp_features/signature_help.zig
index 4979af2a7..f53aabf5c 100644
--- a/tests/lsp_features/signature_help.zig
+++ b/tests/lsp_features/signature_help.zig
@@ -261,7 +261,7 @@ fn testSignatureHelp(source: []const u8, expected_label: []const u8, expected_ac
     var ctx = try Context.init();
     defer ctx.deinit();
 
-    const test_uri = try ctx.addDocument(text);
+    const test_uri = try ctx.addDocument(.{ .source = text });
 
     const params = types.SignatureHelpParams{
         .textDocument = .{ .uri = test_uri },