Add snapshot options, browser_dom_query, and unit test coverage

src/tools/browser/registerBrowserTools.test.ts

@@ -14,6 +14,11 @@ class FakeServer {
 const createSession = (): SessionRecord => {
   const locator = {
     first: vi.fn(() => locator),
+    count: vi.fn(async () => 1),
+    isVisible: vi.fn(async () => true),
+    isEnabled: vi.fn(async () => true),
+    isChecked: vi.fn(async () => false),
+    inputValue: vi.fn(async () => 'test-value'),
     click: vi.fn(async () => undefined),
     fill: vi.fn(async () => undefined),
     hover: vi.fn(async () => undefined),
@@ -30,7 +35,11 @@ const createSession = (): SessionRecord => {
     locator: vi.fn(() => locator),
     waitForSelector: vi.fn(async () => undefined),
     waitForTimeout: vi.fn(async () => undefined),
-    evaluate: vi.fn(async (_callback: unknown, script: string) => `page:${script}`),
+    evaluate: vi.fn(async (_callback: unknown, arg: unknown) => {
+      // extractPageSnapshot passes { maxDepth, maxChildren, selector }; browser_evaluate passes a string
+      if (typeof arg === 'string') return `page:${arg}`;
+      return { tree: [], hiddenTopLevelCount: 0, title: 'Example', url: 'https://example.com' };
+    }),
     keyboard: {
       press: vi.fn(async () => undefined),
       type: vi.fn(async () => undefined),
@@ -172,4 +181,98 @@ describe('registerBrowserTools', () => {
     });
     expect(session.page.mouse.up).toHaveBeenCalled();
   });
+
+  it('browser_snapshot returns a JSON snapshot with params', async () => {
+    const fakeServer = new FakeServer();
+    const session = createSession();
+    const registry = { getSessionOrThrow: vi.fn(() => session) };
+
+    registerBrowserTools(fakeServer as unknown as McpServer, registry as never);
+
+    const snapshot = fakeServer.tools.get('browser_snapshot');
+    const result = (await snapshot?.({
+      sessionId: 1,
+      maxDepth: 3,
+      maxChildren: 10,
+    })) as { content: Array<{ text?: string }> };
+
+    const parsed = JSON.parse(result.content[0]?.text ?? '{}');
+    expect(parsed.title).toBe('Example');
+    expect(parsed.url).toBe('https://example.com');
+    expect(parsed.params).toMatchObject({ maxDepth: 3, maxChildren: 10, selector: null });
+  });
+
+  it('browser_dom_query returns element state when element exists', async () => {
+    const fakeServer = new FakeServer();
+    const session = createSession();
+    const registry = { getSessionOrThrow: vi.fn(() => session) };
+
+    registerBrowserTools(fakeServer as unknown as McpServer, registry as never);
+
+    const domQuery = fakeServer.tools.get('browser_dom_query');
+    const result = (await domQuery?.({
+      sessionId: 1,
+      selector: '#submit',
+    })) as { content: Array<{ text?: string }> };
+
+    const parsed = JSON.parse(result.content[0]?.text ?? '{}');
+    expect(parsed.selector).toBe('#submit');
+    expect(parsed.count).toBe(1);
+    expect(parsed.visible).toBe(true);
+    expect(parsed.enabled).toBe(true);
+  });
+
+  it('browser_dom_query returns only count when element is absent', async () => {
+    const fakeServer = new FakeServer();
+    const session = createSession();
+    // override count to return 0
+    (session.page.locator as ReturnType<typeof vi.fn>).mockReturnValue({
+      count: vi.fn(async () => 0),
+      first: vi.fn(),
+    });
+    const registry = { getSessionOrThrow: vi.fn(() => session) };
+
+    registerBrowserTools(fakeServer as unknown as McpServer, registry as never);
+
+    const domQuery = fakeServer.tools.get('browser_dom_query');
+    const result = (await domQuery?.({
+      sessionId: 1,
+      selector: '.missing',
+    })) as { content: Array<{ text?: string }> };
+
+    const parsed = JSON.parse(result.content[0]?.text ?? '{}');
+    expect(parsed.selector).toBe('.missing');
+    expect(parsed.count).toBe(0);
+    expect(parsed.visible).toBeUndefined();
+  });
+
+  it('browser_dom_query sets checked and value to null for non-checkbox non-input elements', async () => {
+    const fakeServer = new FakeServer();
+    const session = createSession();
+    // isChecked and inputValue throw for non-checkbox / non-input elements
+    const firstLocator = {
+      isVisible: vi.fn(async () => true),
+      isEnabled: vi.fn(async () => true),
+      isChecked: vi.fn(async () => { throw new Error('not a checkbox'); }),
+      inputValue: vi.fn(async () => { throw new Error('not an input'); }),
+    };
+    (session.page.locator as ReturnType<typeof vi.fn>).mockReturnValue({
+      count: vi.fn(async () => 1),
+      first: vi.fn(() => firstLocator),
+    });
+    const registry = { getSessionOrThrow: vi.fn(() => session) };
+
+    registerBrowserTools(fakeServer as unknown as McpServer, registry as never);
+
+    const domQuery = fakeServer.tools.get('browser_dom_query');
+    const result = (await domQuery?.({
+      sessionId: 1,
+      selector: 'span',
+    })) as { content: Array<{ text?: string }> };
+
+    const parsed = JSON.parse(result.content[0]?.text ?? '{}');
+    expect(parsed.count).toBe(1);
+    expect(parsed.checked).toBeNull();
+    expect(parsed.value).toBeNull();
+  });
 });

src/tools/browser/registerBrowserTools.ts

@@ -1,6 +1,7 @@
 import type { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
 import { SessionRegistry } from '../../sessions/SessionRegistry.js';
 import {
+  domQuerySchema,
   dragSchema,
   evaluateSchema,
   fillFormSchema,
@@ -15,6 +16,7 @@ import {
   selectOptionSchema,
   selectorSchema,
   sessionIdSchema,
+  snapshotSchema,
   uploadFileSchema,
   waitForSelectorSchema,
   waitForTimeoutSchema,
@@ -127,10 +129,10 @@ export const registerBrowserTools = (server: McpServer, registry: SessionRegistr
     {
       title: 'Browser Snapshot',
       description: 'Return a structured page snapshot.',
-      inputSchema: sessionIdSchema,
+      inputSchema: snapshotSchema,
     },
-    withSession(registry, async ({ page }) => {
-      const snapshot = await extractPageSnapshot(page);
+    withSession(registry, async ({ page }, { maxDepth, maxChildren, selector }) => {
+      const snapshot = await extractPageSnapshot(page, { maxDepth, maxChildren, selector });
 
       return jsonResult(snapshot);
     }),
@@ -210,15 +212,51 @@ export const registerBrowserTools = (server: McpServer, registry: SessionRegistr
     {
       title: 'Browser Get Page Structure',
       description: 'Return a readable page structure summary.',
-      inputSchema: sessionIdSchema,
+      inputSchema: snapshotSchema,
     },
-    withSession(registry, async ({ page }) => {
-      const snapshot = await extractPageSnapshot(page);
+    withSession(registry, async ({ page }, { maxDepth, maxChildren, selector }) => {
+      const snapshot = await extractPageSnapshot(page, { maxDepth, maxChildren, selector });
 
       return textResult(formatPageStructure(snapshot));
     }),
   );
 
+  server.registerTool(
+    'browser_dom_query',
+    {
+      title: 'Browser DOM Query',
+      description: 'Query element presence, count, and state without waiting.',
+      inputSchema: domQuerySchema,
+    },
+    withSession(registry, async ({ page }, { selector }) => {
+      const locator = page.locator(selector);
+      const count = await locator.count();
+
+      if (count === 0) {
+        return jsonResult({ selector, count });
+      }
+
+      const first = locator.first();
+      const [visible, enabled] = await Promise.all([first.isVisible(), first.isEnabled()]);
+
+      let checked: boolean | null = null;
+      try {
+        checked = await first.isChecked();
+      } catch {
+        // not a checkbox or radio
+      }
+
+      let value: string | null = null;
+      try {
+        value = await first.inputValue();
+      } catch {
+        // not an input element
+      }
+
+      return jsonResult({ selector, count, visible, enabled, checked, value });
+    }),
+  );
+
   server.registerTool(
     'browser_evaluate',
     {

src/types/toolArgs.ts

@@ -37,7 +37,17 @@ export const fillFormSchema = sessionIdSchema.extend({
 });
 
 export const screenshotSchema = sessionIdSchema.extend({
-  fullPage: z.boolean().optional(),
+  fullPage: z.boolean().optional().describe('Capture full scrollable page, not just the viewport.'),
+});
+
+export const snapshotSchema = sessionIdSchema.extend({
+  maxDepth: z.number().int().positive().optional().describe('Max tree depth (default 4).'),
+  maxChildren: z.number().int().positive().optional().describe('Max children per node (default 20).'),
+  selector: z.string().min(1).optional().describe('Scope to an element; omit for document body.'),
+});
+
+export const domQuerySchema = sessionIdSchema.extend({
+  selector: z.string().min(1),
 });
 
 export const dragSchema = sessionIdSchema.extend({
@@ -48,7 +58,7 @@ export const dragSchema = sessionIdSchema.extend({
 
 export const selectOptionSchema = sessionIdSchema.extend({
   selector: z.string().min(1),
-  values: z.array(z.string().min(1)).min(1),
+  values: z.array(z.string().min(1)).min(1).describe('Option values to select; pass multiple for multi-select.'),
   timeout: z.number().int().positive().optional(),
 });
 
@@ -58,16 +68,16 @@ export const generateLocatorSchema = sessionIdSchema.extend({
 
 export const evaluateSchema = sessionIdSchema.extend({
   script: z.string().min(1),
-  selector: z.string().min(1).optional(),
+  selector: z.string().min(1).optional().describe('Matched element passed as first arg to script.'),
 });
 
 export const keyboardPressSchema = sessionIdSchema.extend({
-  key: z.string().min(1),
+  key: z.string().min(1).describe('Playwright key name, e.g. "Enter", "ArrowDown", "Control+A".'),
 });
 
 export const keyboardTypeSchema = sessionIdSchema.extend({
   text: z.string(),
-  delay: z.number().int().nonnegative().optional(),
+  delay: z.number().int().nonnegative().optional().describe('Ms between keystrokes. Omit for instant.'),
 });
 
 export const mousePointSchema = sessionIdSchema.extend({