From 6008187c78c33824dec793b6dd6414419f3ad4ed Mon Sep 17 00:00:00 2001
From: Matt Van Horn <455140+mvanhorn@users.noreply.github.com>
Date: Fri, 20 Mar 2026 21:38:11 -0700
Subject: [PATCH] Add waitForSelector MCP tool

Adds a waitForSelector tool to the MCP server that polls for a CSS
selector match with a configurable timeout (default 5000ms). Returns the
backendNodeId of the matched element for use with click/fill tools.

The tool runs the session event loop between selector checks, so
dynamically-created elements are found as they appear from JS execution
or network responses.
---
 src/mcp/tools.zig | 53 +++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 53 insertions(+)

diff --git a/src/mcp/tools.zig b/src/mcp/tools.zig
index 51a44476..48ba3de5 100644
--- a/src/mcp/tools.zig
+++ b/src/mcp/tools.zig
@@ -142,6 +142,20 @@ pub const tool_list = [_]protocol.Tool{
             \\}
         ),
     },
+    .{
+        .name = "waitForSelector",
+        .description = "Wait for an element matching a CSS selector to appear in the page. Returns the backend node ID of the matched element.",
+        .inputSchema = protocol.minify(
+            \\{
+            \\  "type": "object",
+            \\  "properties": {
+            \\    "selector": { "type": "string", "description": "The CSS selector to wait for." },
+            \\    "timeout": { "type": "integer", "description": "Optional timeout in milliseconds. Defaults to 5000." }
+            \\  },
+            \\  "required": ["selector"]
+            \\}
+        ),
+    },
 };
 
 pub fn handleList(server: *Server, arena: std.mem.Allocator, req: protocol.Request) !void {
@@ -243,6 +257,7 @@ const ToolAction = enum {
     click,
     fill,
     scroll,
+    waitForSelector,
 };
 
 const tool_map = std.StaticStringMap(ToolAction).initComptime(.{
@@ -257,6 +272,7 @@ const tool_map = std.StaticStringMap(ToolAction).initComptime(.{
     .{ "click", .click },
     .{ "fill", .fill },
     .{ "scroll", .scroll },
+    .{ "waitForSelector", .waitForSelector },
 });
 
 pub fn handleCall(server: *Server, arena: std.mem.Allocator, req: protocol.Request) !void {
@@ -288,6 +304,7 @@ pub fn handleCall(server: *Server, arena: std.mem.Allocator, req: protocol.Reque
         .click => try handleClick(server, arena, req.id.?, call_params.arguments),
         .fill => try handleFill(server, arena, req.id.?, call_params.arguments),
         .scroll => try handleScroll(server, arena, req.id.?, call_params.arguments),
+        .waitForSelector => try handleWaitForSelector(server, arena, req.id.?, call_params.arguments),
     }
 }
 
@@ -532,6 +549,42 @@ fn handleScroll(server: *Server, arena: std.mem.Allocator, id: std.json.Value, a
     const content = [_]protocol.TextContent([]const u8){.{ .text = "Scrolled successfully." }};
     try server.sendResult(id, protocol.CallToolResult([]const u8){ .content = &content });
 }
+fn handleWaitForSelector(server: *Server, arena: std.mem.Allocator, id: std.json.Value, arguments: ?std.json.Value) !void {
+    const WaitParams = struct {
+        selector: [:0]const u8,
+        timeout: ?u32 = null,
+    };
+    const args = try parseArguments(WaitParams, arena, arguments, server, id, "waitForSelector");
+
+    const page = server.session.currentPage() orelse {
+        return server.sendError(id, .PageNotLoaded, "Page not loaded");
+    };
+
+    const timeout_ms = args.timeout orelse 5000;
+    var timer = try std.time.Timer.start();
+
+    while (true) {
+        const element = Selector.querySelector(page.document.asNode(), args.selector, page) catch {
+            return server.sendError(id, .InvalidParams, "Invalid selector");
+        };
+
+        if (element) |el| {
+            const registered = try server.node_registry.register(el.asNode());
+            const msg = std.fmt.allocPrint(arena, "Element found. backendNodeId: {d}", .{registered.id}) catch "Element found.";
+
+            const content = [_]protocol.TextContent([]const u8){.{ .text = msg }};
+            return server.sendResult(id, protocol.CallToolResult([]const u8){ .content = &content });
+        }
+
+        const elapsed: u32 = @intCast(timer.read() / std.time.ns_per_ms);
+        if (elapsed >= timeout_ms) {
+            return server.sendError(id, .InternalError, "Timeout waiting for selector");
+        }
+
+        _ = server.session.wait(.{ .timeout_ms = @min(100, timeout_ms - elapsed) });
+    }
+}
+
 fn parseArguments(comptime T: type, arena: std.mem.Allocator, arguments: ?std.json.Value, server: *Server, id: std.json.Value, tool_name: []const u8) !T {
     if (arguments == null) {
         try server.sendError(id, .InvalidParams, "Missing arguments");