From 814ca8ab3fdfa69e706622aede221184cc29f7fb Mon Sep 17 00:00:00 2001
From: Navid EMAD <navid.emad@yespark.fr>
Date: Tue, 19 May 2026 20:43:54 +0200
Subject: [PATCH] accessibility: unify query+tree writers, route objectId via
 dom.getNode

Fold QueryWriter into Writer behind an Opts.filter. Tree mode is unchanged
(filter=null); query mode walks the full subtree (including AX-ignored
nodes per the queryAXTree spec) and emits the flat-match shape. Shared
resolveRole helper handles label-promotion for both paths so the two
can't drift.

Drop the "objectId not yet supported" carve-out: queryAXTree now reuses
dom.getNode, which already resolves nodeId/backendNodeId/objectId.
---
 src/cdp/AXNode.zig                | 176 ++++++++++++++----------------
 src/cdp/CDP.zig                   |  22 +---
 src/cdp/domains/accessibility.zig |  46 ++------
 src/cdp/domains/dom.zig           |   2 +-
 4 files changed, 92 insertions(+), 154 deletions(-)

diff --git a/src/cdp/AXNode.zig b/src/cdp/AXNode.zig
index dd781b99..1621085d 100644
--- a/src/cdp/AXNode.zig
+++ b/src/cdp/AXNode.zig
@@ -45,11 +45,31 @@ pub const Writer = struct {
     visibility_cache: *DOMNode.Element.VisibilityCache,
     label_index: *Label.LabelByForIndex,
     temp_arena: std.mem.Allocator,
+    // When null, emit the full AX tree (getFullAXTree). When set, walk the
+    // subtree visiting all nodes (including AX-ignored ones, per the
+    // queryAXTree spec) and emit only nodes whose role + accessible name
+    // match the filter, in a flat shape.
+    filter: ?Filter = null,
 
-    pub const Opts = struct {};
+    pub const Filter = struct {
+        role: ?[]const u8 = null,
+        accessible_name: ?[]const u8 = null,
+    };
+
+    pub const Opts = struct {
+        filter: ?Filter = null,
+    };
+
+    const ResolvedRole = struct {
+        role: []const u8,
+        // Non-null when the current node is a <label> whose target control is
+        // a CSS-hidden checkbox/radio. Callers use this to emit the control's
+        // checked/disabled properties on the promoted label.
+        promoted_input: ?*DOMNode.Element.Html.Input,
+    };
 
     pub fn jsonStringify(self: *const Writer, w: anytype) error{WriteFailed}!void {
-        self.toJSON(self.root, w) catch |err| {
+        self.toJSON(w) catch |err| {
             // The only error our jsonStringify method can return is
             // @TypeOf(w).Error. In other words, our code can't return its own
             // error, we can only return a writer error. Kinda sucks.
@@ -58,15 +78,36 @@ pub const Writer = struct {
         };
     }
 
-    fn toJSON(self: *const Writer, node: *const Node, w: anytype) !void {
+    fn toJSON(self: *const Writer, w: anytype) !void {
         try w.beginArray();
-        const root = AXNode.fromNode(node.dom);
-        if (try self.writeNode(node.id, root, false, w)) {
-            try self.writeNodeChildren(root, false, w);
+        if (self.filter != null) {
+            try self.walkQuery(self.root.dom, false, w);
+        } else {
+            const root = AXNode.fromNode(self.root.dom);
+            if (try self.writeNode(self.root.id, root, false, w)) {
+                try self.writeNodeChildren(root, false, w);
+            }
         }
         return w.endArray();
     }
 
+    // Resolve the displayed role for `axn`, accounting for label-promotion
+    // when a <label> targets a CSS-hidden checkbox/radio. Shared between the
+    // tree (writeNode) and query (emitMatch) paths so the two can't drift.
+    fn resolveRole(self: *const Writer, axn: AXNode) !ResolvedRole {
+        if (labelPromotionTarget(axn, self.frame, self.visibility_cache)) |input| {
+            return .{
+                .role = switch (input._input_type) {
+                    .checkbox => "checkbox",
+                    .radio => "radio",
+                    else => unreachable,
+                },
+                .promoted_input = input,
+            };
+        }
+        return .{ .role = try axn.getRole(), .promoted_input = null };
+    }
+
     // CDP spec defines AXNodeId as a string, so nodeId/parentId/childIds must
     // be serialized as JSON strings even though we track them internally as u32.
     fn writeIdString(id: u32, w: anytype) !void {
@@ -485,18 +526,11 @@ pub const Writer = struct {
         try w.objectField("backendDOMNodeId");
         try w.write(id);
 
-        const promoted_input = labelPromotionTarget(axn, self.frame, self.visibility_cache);
+        const resolved = try self.resolveRole(axn);
+        const promoted_input = resolved.promoted_input;
 
         try w.objectField("role");
-        if (promoted_input) |input| {
-            try self.writeAXValue(.{ .role = switch (input._input_type) {
-                .checkbox => "checkbox",
-                .radio => "radio",
-                else => unreachable,
-            } }, w);
-        } else {
-            try self.writeAXValue(.{ .role = try axn.getRole() }, w);
-        }
+        try self.writeAXValue(.{ .role = resolved.role }, w);
 
         const ignore = axn.isIgnore(self.frame, self.visibility_cache, in_aria_hidden);
         try w.objectField("ignored");
@@ -632,47 +666,12 @@ pub const Writer = struct {
             try self.writeAXValue(.{ .string = val }, w);
         }
     }
-};
 
-// Backing writer for Accessibility.queryAXTree. Walks the DOM subtree rooted
-// at `root`, computes role + accessible name for every node (including those
-// that would be ignored in the AX tree, per the CDP spec), and emits a flat
-// JSON array of nodes whose role/name match the optional filters.
-//
-// Intentionally MVP: emits an empty `properties` array and empty `childIds`.
-// Clients that need full properties call Accessibility.getFullAXTree on a
-// matched nodeId.
-pub const QueryWriter = struct {
-    root: *const Node,
-    registry: *Node.Registry,
-    frame: *Frame,
-    visibility_cache: *DOMNode.Element.VisibilityCache,
-    label_index: *Label.LabelByForIndex,
-    temp_arena: std.mem.Allocator,
-    accessible_name: ?[]const u8,
-    role: ?[]const u8,
-
-    pub const Opts = struct {
-        accessible_name: ?[]const u8 = null,
-        role: ?[]const u8 = null,
-    };
-
-    pub fn jsonStringify(self: *const QueryWriter, w: anytype) error{WriteFailed}!void {
-        self.toJSON(w) catch |err| {
-            log.err(.cdp, "queryAXTree toJSON", .{ .err = err });
-            return error.WriteFailed;
-        };
-    }
-
-    fn toJSON(self: *const QueryWriter, w: anytype) !void {
-        try w.beginArray();
-        try self.walk(self.root.dom, false, w);
-        return w.endArray();
-    }
-
-    fn walk(self: *const QueryWriter, node: *DOMNode, in_aria_hidden: bool, w: anytype) !void {
+    // Query-mode walk. Visits every node under `root` (including AX-ignored
+    // ones, per the queryAXTree spec) and defers emission to emitMatch.
+    fn walkQuery(self: *const Writer, node: *DOMNode, in_aria_hidden: bool, w: anytype) !void {
         const axn = AXNode.fromNode(node);
-        try self.maybeEmit(axn, in_aria_hidden, w);
+        try self.emitMatch(axn, in_aria_hidden, w);
 
         // <head>, <script>, <style> never expose AX content — skip recursion.
         if (axn.ignoreChildren()) return;
@@ -688,30 +687,24 @@ pub const QueryWriter = struct {
                 .element, .cdata => {},
                 else => continue,
             }
-            try self.walk(child, child_in_aria_hidden, w);
+            try self.walkQuery(child, child_in_aria_hidden, w);
         }
     }
 
-    fn maybeEmit(self: *const QueryWriter, axn: AXNode, in_aria_hidden: bool, w: anytype) !void {
-        // Mirror Writer.writeNode (line 488-499): if this is a <label> for a
-        // hidden checkbox/radio (toggle-switch / CSS-only radio pattern), emit
-        // the label with the input's role so clients searching by role land
-        // on the label's clickable backendDOMNodeId rather than the hidden,
-        // non-interactable input.
-        const promoted_input = labelPromotionTarget(axn, self.frame, self.visibility_cache);
+    // Emit `axn` to `w` iff it satisfies the active filter. Output shape is
+    // the queryAXTree flat-match shape: nodeId, backendDOMNodeId, ignored,
+    // role, name, plus empty properties / childIds (clients fetch full
+    // properties via getFullAXTree on a matched nodeId).
+    fn emitMatch(self: *const Writer, axn: AXNode, in_aria_hidden: bool, w: anytype) !void {
+        const filter = self.filter.?;
+        const resolved = self.resolveRole(axn) catch return;
 
-        const role_str = if (promoted_input) |input| switch (input._input_type) {
-            .checkbox => "checkbox",
-            .radio => "radio",
-            else => unreachable,
-        } else (axn.getRole() catch return);
-
-        if (self.role) |needle| {
-            if (!std.mem.eql(u8, needle, role_str)) return;
+        if (filter.role) |needle| {
+            if (!std.mem.eql(u8, needle, resolved.role)) return;
         }
 
         const name = (try axn.getName(self.frame, self.temp_arena)) orelse "";
-        if (self.accessible_name) |needle| {
+        if (filter.accessible_name) |needle| {
             if (!std.mem.eql(u8, needle, name)) return;
         }
 
@@ -721,9 +714,7 @@ pub const QueryWriter = struct {
         try w.beginObject();
 
         try w.objectField("nodeId");
-        var buf: [10]u8 = undefined;
-        const id_str = try std.fmt.bufPrint(&buf, "{d}", .{node.id});
-        try w.write(id_str);
+        try writeIdString(node.id, w);
 
         try w.objectField("backendDOMNodeId");
         try w.write(node.id);
@@ -732,12 +723,7 @@ pub const QueryWriter = struct {
         try w.write(ignored);
 
         try w.objectField("role");
-        try w.beginObject();
-        try w.objectField("type");
-        try w.write("role");
-        try w.objectField("value");
-        try w.write(role_str);
-        try w.endObject();
+        try self.writeAXValue(.{ .role = resolved.role }, w);
 
         try w.objectField("name");
         try w.beginObject();
@@ -1717,7 +1703,7 @@ test "AXNode: writer prunes hidden and resolves labels" {
     }
 }
 
-test "AXNode: QueryWriter filters by role" {
+test "AXNode: Writer query filters by role" {
     var registry = Node.Registry.init(testing.allocator);
     defer registry.deinit();
 
@@ -1731,15 +1717,14 @@ test "AXNode: QueryWriter filters by role" {
     const temp_arena = try frame.getArena(.medium, "AXNode");
     defer frame.releaseArena(temp_arena);
 
-    const json = try std.json.Stringify.valueAlloc(testing.allocator, QueryWriter{
+    const json = try std.json.Stringify.valueAlloc(testing.allocator, Writer{
         .root = node,
         .registry = &registry,
         .frame = frame,
         .visibility_cache = &visibility_cache,
         .label_index = &label_index,
         .temp_arena = temp_arena,
-        .accessible_name = null,
-        .role = "heading",
+        .filter = .{ .role = "heading" },
     }, .{});
     defer testing.allocator.free(json);
 
@@ -1757,7 +1742,7 @@ test "AXNode: QueryWriter filters by role" {
     try testing.expectEqual("Visible", name_val);
 }
 
-test "AXNode: QueryWriter filters by accessible name" {
+test "AXNode: Writer query filters by accessible name" {
     var registry = Node.Registry.init(testing.allocator);
     defer registry.deinit();
 
@@ -1771,15 +1756,14 @@ test "AXNode: QueryWriter filters by accessible name" {
     const temp_arena = try frame.getArena(.medium, "AXNode");
     defer frame.releaseArena(temp_arena);
 
-    const json = try std.json.Stringify.valueAlloc(testing.allocator, QueryWriter{
+    const json = try std.json.Stringify.valueAlloc(testing.allocator, Writer{
         .root = node,
         .registry = &registry,
         .frame = frame,
         .visibility_cache = &visibility_cache,
         .label_index = &label_index,
         .temp_arena = temp_arena,
-        .accessible_name = "Search",
-        .role = null,
+        .filter = .{ .accessible_name = "Search" },
     }, .{});
     defer testing.allocator.free(json);
 
@@ -1797,7 +1781,7 @@ test "AXNode: QueryWriter filters by accessible name" {
     }
 }
 
-test "AXNode: QueryWriter combined role and name filter promotes hidden-input labels" {
+test "AXNode: Writer query combined role+name filter promotes hidden-input labels" {
     var registry = Node.Registry.init(testing.allocator);
     defer registry.deinit();
 
@@ -1817,15 +1801,14 @@ test "AXNode: QueryWriter combined role and name filter promotes hidden-input la
     //   - the hidden input (intrinsic role="checkbox", ignored=true)
     // The label is the actionable target — a real client searching by role
     // would act on the entry with ignored=false.
-    const json = try std.json.Stringify.valueAlloc(testing.allocator, QueryWriter{
+    const json = try std.json.Stringify.valueAlloc(testing.allocator, Writer{
         .root = node,
         .registry = &registry,
         .frame = frame,
         .visibility_cache = &visibility_cache,
         .label_index = &label_index,
         .temp_arena = temp_arena,
-        .accessible_name = "Enable feature",
-        .role = "checkbox",
+        .filter = .{ .accessible_name = "Enable feature", .role = "checkbox" },
     }, .{});
     defer testing.allocator.free(json);
 
@@ -1849,7 +1832,7 @@ test "AXNode: QueryWriter combined role and name filter promotes hidden-input la
     try testing.expect(actionable_match);
 }
 
-test "AXNode: QueryWriter no match returns empty array" {
+test "AXNode: Writer query no match returns empty array" {
     var registry = Node.Registry.init(testing.allocator);
     defer registry.deinit();
 
@@ -1863,15 +1846,14 @@ test "AXNode: QueryWriter no match returns empty array" {
     const temp_arena = try frame.getArena(.medium, "AXNode");
     defer frame.releaseArena(temp_arena);
 
-    const json = try std.json.Stringify.valueAlloc(testing.allocator, QueryWriter{
+    const json = try std.json.Stringify.valueAlloc(testing.allocator, Writer{
         .root = node,
         .registry = &registry,
         .frame = frame,
         .visibility_cache = &visibility_cache,
         .label_index = &label_index,
         .temp_arena = temp_arena,
-        .accessible_name = null,
-        .role = "marquee",
+        .filter = .{ .role = "marquee" },
     }, .{});
     defer testing.allocator.free(json);
 
diff --git a/src/cdp/CDP.zig b/src/cdp/CDP.zig
index e44e1708..f23e8f8a 100644
--- a/src/cdp/CDP.zig
+++ b/src/cdp/CDP.zig
@@ -661,7 +661,6 @@ pub const BrowserContext = struct {
         // names/visibility from the wrong document.
         const fallback = self.session.currentFrame() orelse return error.FrameNotLoaded;
         const frame = root.dom.ownerFrame(fallback);
-        _ = opts;
         const cache = try frame.call_arena.create(Element.VisibilityCache);
         cache.* = .empty;
         const label_index = try frame.call_arena.create(Label.LabelByForIndex);
@@ -673,26 +672,7 @@ pub const BrowserContext = struct {
             .visibility_cache = cache,
             .label_index = label_index,
             .temp_arena = temp_arena,
-        };
-    }
-
-    pub fn axnodeQueryWriter(self: *BrowserContext, temp_arena: Allocator, root: *const Node, opts: AXNode.QueryWriter.Opts) !AXNode.QueryWriter {
-        // See axnodeWriter for the frame-binding rationale.
-        const fallback = self.session.currentFrame() orelse return error.FrameNotLoaded;
-        const frame = root.dom.ownerFrame(fallback);
-        const cache = try frame.call_arena.create(Element.VisibilityCache);
-        cache.* = .empty;
-        const label_index = try frame.call_arena.create(Label.LabelByForIndex);
-        label_index.* = .{};
-        return .{
-            .root = root,
-            .registry = &self.node_registry,
-            .frame = frame,
-            .visibility_cache = cache,
-            .label_index = label_index,
-            .temp_arena = temp_arena,
-            .accessible_name = opts.accessible_name,
-            .role = opts.role,
+            .filter = opts.filter,
         };
     }
 
diff --git a/src/cdp/domains/accessibility.zig b/src/cdp/domains/accessibility.zig
index 12f906b0..1ff28aae 100644
--- a/src/cdp/domains/accessibility.zig
+++ b/src/cdp/domains/accessibility.zig
@@ -19,6 +19,7 @@
 const std = @import("std");
 const id = @import("../id.zig");
 const CDP = @import("../CDP.zig");
+const dom = @import("dom.zig");
 
 pub fn processMessage(cmd: *CDP.Command) !void {
     const action = std.meta.stringToEnum(enum {
@@ -78,63 +79,38 @@ fn queryAXTree(cmd: *CDP.Command) !void {
         accessibleName: ?[]const u8 = null,
         role: ?[]const u8 = null,
     };
-    // Default-construct on missing params so we can return our specific
-    // "node identifier required" error rather than a generic InvalidParams.
-    const params = (try cmd.params(Params)) orelse Params{};
-
-    // objectId requires the JS inspector and an attached runtime — defer to a
-    // follow-up. Real clients (Capybara, Stagehand) usually have a nodeId from
-    // a prior DOM.querySelector/DOM.getDocument call.
-    if (params.objectId != null and params.nodeId == null and params.backendNodeId == null) {
-        return cmd.sendError(-32000, "Accessibility.queryAXTree by objectId is not yet supported; use nodeId or backendNodeId", .{});
-    }
-
-    const input_id = params.nodeId orelse params.backendNodeId orelse {
-        return cmd.sendError(-32000, "Either nodeId, backendNodeId or objectId must be specified", .{});
-    };
+    const params = (try cmd.params(Params)) orelse return error.InvalidParams;
 
     const bc = cmd.browser_context orelse return error.BrowserContextNotLoaded;
-    const node = bc.node_registry.lookup_by_id.get(input_id) orelse return error.NodeNotFound;
+    const node = try dom.getNode(cmd.arena, bc, params.nodeId, params.backendNodeId, params.objectId);
 
     const frame = bc.session.currentFrame() orelse return error.FrameNotLoaded;
     const temp_arena = try frame.getArena(.medium, "AXNode");
     defer frame.releaseArena(temp_arena);
 
-    return cmd.sendResult(.{ .nodes = try bc.axnodeQueryWriter(temp_arena, node, .{
-        .accessible_name = params.accessibleName,
-        .role = params.role,
+    return cmd.sendResult(.{ .nodes = try bc.axnodeWriter(temp_arena, node, .{
+        .filter = .{
+            .accessible_name = params.accessibleName,
+            .role = params.role,
+        },
     }) }, .{});
 }
 
 const testing = @import("../testing.zig");
 
-test "cdp.accessibility: queryAXTree requires nodeId or backendNodeId" {
+test "cdp.accessibility: queryAXTree requires nodeId, backendNodeId or objectId" {
     var ctx = try testing.context();
     defer ctx.deinit();
 
     _ = try ctx.loadBrowserContext(.{ .id = "BID-A", .url = "cdp/ax_tree.html" });
 
-    // Pass filters but no node identifier — exercises the missing-id branch.
+    // Pass filters but no node identifier — dom.getNode returns MissingParams.
     try ctx.processMessage(.{
         .id = 1,
         .method = "Accessibility.queryAXTree",
         .params = .{ .role = "button" },
     });
-    try ctx.expectSentError(-32000, "Either nodeId, backendNodeId or objectId must be specified", .{ .id = 1 });
-}
-
-test "cdp.accessibility: queryAXTree with objectId only is not yet supported" {
-    var ctx = try testing.context();
-    defer ctx.deinit();
-
-    _ = try ctx.loadBrowserContext(.{ .id = "BID-A", .url = "cdp/ax_tree.html" });
-
-    try ctx.processMessage(.{
-        .id = 1,
-        .method = "Accessibility.queryAXTree",
-        .params = .{ .objectId = "OBJ-X" },
-    });
-    try ctx.expectSentError(-32000, "Accessibility.queryAXTree by objectId is not yet supported; use nodeId or backendNodeId", .{ .id = 1 });
+    try ctx.expectSentError(-31998, "MissingParams", .{ .id = 1 });
 }
 
 test "cdp.accessibility: queryAXTree with unknown nodeId returns error" {
diff --git a/src/cdp/domains/dom.zig b/src/cdp/domains/dom.zig
index aa7cfb2a..a164e79f 100644
--- a/src/cdp/domains/dom.zig
+++ b/src/cdp/domains/dom.zig
@@ -459,7 +459,7 @@ fn scrollIntoViewIfNeeded(cmd: *CDP.Command) !void {
     return cmd.sendResult(null, .{});
 }
 
-fn getNode(arena: Allocator, bc: *CDP.BrowserContext, node_id: ?Node.Id, backend_node_id: ?Node.Id, object_id: ?[]const u8) !*Node {
+pub fn getNode(arena: Allocator, bc: *CDP.BrowserContext, node_id: ?Node.Id, backend_node_id: ?Node.Id, object_id: ?[]const u8) !*Node {
     const input_node_id = node_id orelse backend_node_id;
     if (input_node_id) |input_node_id_| {
         return bc.node_registry.lookup_by_id.get(input_node_id_) orelse return error.NodeNotFound;