feat(tools): integrate expectation parameter into main tools

- Add expectation parameter to navigate, click, type, and snapshot tools
- Configure tool-specific default expectations for optimal token usage
- Update browserServerBackend to pass expectation to Response class
- Remove legacy setIncludeSnapshot/setIncludeTabs methods
- Ensure complete backward compatibility

This completes PR #3 of the token optimization implementation,
enabling fine-grained control over response content per tool.

Author: Tomohiko Hiraki

src/browserServerBackend.ts

@@ -77,7 +77,7 @@ export class BrowserServerBackend implements ServerBackend {
 
   async callTool(schema: mcpServer.ToolSchema<any>, parsedArguments: any) {
     const context = this._context!;
-    const response = new Response(context, schema.name, parsedArguments);
+    const response = new Response(context, schema.name, parsedArguments, parsedArguments.expectation);
     const tool = this._tools.find(tool => tool.schema.name === schema.name)!;
     context.setRunningTool(true);
     try {

src/response.ts

@@ -87,7 +87,9 @@ export class Response {
   async finish() {
     // All the async snapshotting post-action is happening here.
     // Everything below should race against modal states.
-    if ((this._includeSnapshot || this._expectation.includeSnapshot) && this._context.currentTab()) {
+    // Expectation settings take priority over legacy setIncludeSnapshot calls
+    const shouldIncludeSnapshot = this._expectation.includeSnapshot;
+    if (shouldIncludeSnapshot && this._context.currentTab()) {
       const options = this._expectation.snapshotOptions;
       if (options?.selector) {
         // TODO: Implement partial snapshot capture based on selector
@@ -125,10 +127,10 @@ ${this._code.join('\n')}
     }
 
     // List browser tabs based on expectation.
-    const shouldIncludeTabs = this._expectation.includeTabs || this._includeTabs;
-    const shouldIncludeSnapshot = this._expectation.includeSnapshot || this._includeSnapshot;
+    const shouldIncludeTabs = this._expectation.includeTabs;
+    const shouldIncludeSnapshot = this._expectation.includeSnapshot;
 
-    if (shouldIncludeSnapshot || shouldIncludeTabs)
+    if (shouldIncludeTabs)
       response.push(...renderTabsMarkdown(this._context.tabs(), shouldIncludeTabs));
 
     // Add snapshot if provided and expectation allows it.

src/schemas/expectation.ts

@@ -52,7 +52,7 @@ export type ExpectationOptions = z.infer<typeof expectationSchema>;
  */
 const TOOL_DEFAULTS: Record<string, Required<Omit<NonNullable<ExpectationOptions>, 'snapshotOptions' | 'consoleOptions' | 'imageOptions'>>> = {
   // Navigation tools need full context for verification
-  navigate: {
+  browser_navigate: {
     includeSnapshot: true,
     includeConsole: true,
     includeDownloads: true,
@@ -61,15 +61,15 @@ const TOOL_DEFAULTS: Record<string, Required<Omit<NonNullable<ExpectationOptions
   },
 
   // Interactive tools need snapshot for feedback but less verbose logging
-  click: {
+  browser_click: {
     includeSnapshot: true,
     includeConsole: false,
     includeDownloads: false,
     includeTabs: false,
     includeCode: true
   },
 
-  fill: {
+  browser_type: {
     includeSnapshot: true,
     includeConsole: false,
     includeDownloads: false,
@@ -86,6 +86,15 @@ const TOOL_DEFAULTS: Record<string, Required<Omit<NonNullable<ExpectationOptions
     includeCode: false
   },
 
+  // Snapshot tool should capture snapshot but minimal other context
+  browser_snapshot: {
+    includeSnapshot: true,
+    includeConsole: false,
+    includeDownloads: false,
+    includeTabs: false,
+    includeCode: false
+  },
+  
   // Code evaluation needs console output but minimal other info
   evaluate: {
     includeSnapshot: false,

src/tools/keyboard.ts

@@ -20,6 +20,7 @@ import { defineTabTool } from './tool.js';
 import { elementSchema } from './snapshot.js';
 import { generateLocator } from './utils.js';
 import * as javascript from '../javascript.js';
+import { expectationSchema } from '../schemas/expectation.js';
 
 const pressKey = defineTabTool({
   capability: 'core',
@@ -30,12 +31,12 @@ const pressKey = defineTabTool({
     description: 'Press a key on the keyboard',
     inputSchema: z.object({
       key: z.string().describe('Name of the key to press or a character to generate, such as `ArrowLeft` or `a`'),
+      expectation: expectationSchema
     }),
     type: 'destructive',
   },
 
   handle: async (tab, params, response) => {
-    response.setIncludeSnapshot();
     response.addCode(`// Press ${params.key}`);
     response.addCode(`await page.keyboard.press('${params.key}');`);
 
@@ -49,6 +50,7 @@ const typeSchema = elementSchema.extend({
   text: z.string().describe('Text to type into the element'),
   submit: z.boolean().optional().describe('Whether to submit entered text (press Enter after)'),
   slowly: z.boolean().optional().describe('Whether to type one character at a time. Useful for triggering key handlers in the page. By default entire text is filled in at once.'),
+  expectation: expectationSchema
 });
 
 const type = defineTabTool({
@@ -66,7 +68,6 @@ const type = defineTabTool({
 
     await tab.waitForCompletion(async () => {
       if (params.slowly) {
-        response.setIncludeSnapshot();
         response.addCode(`await page.${await generateLocator(locator)}.pressSequentially(${javascript.quote(params.text)});`);
         await locator.pressSequentially(params.text);
       } else {
@@ -75,7 +76,6 @@ const type = defineTabTool({
       }
 
       if (params.submit) {
-        response.setIncludeSnapshot();
         response.addCode(`await page.${await generateLocator(locator)}.press('Enter');`);
         await locator.press('Enter');
       }

src/tools/navigate.ts

@@ -16,6 +16,7 @@
 
 import { z } from 'zod';
 import { defineTool, defineTabTool } from './tool.js';
+import { expectationSchema } from '../schemas/expectation.js';
 
 const navigate = defineTool({
   capability: 'core',
@@ -26,6 +27,7 @@ const navigate = defineTool({
     description: 'Navigate to a URL',
     inputSchema: z.object({
       url: z.string().describe('The URL to navigate to'),
+      expectation: expectationSchema
     }),
     type: 'destructive',
   },
@@ -34,7 +36,6 @@ const navigate = defineTool({
     const tab = await context.ensureTab();
     await tab.navigate(params.url);
 
-    response.setIncludeSnapshot();
     response.addCode(`await page.goto('${params.url}');`);
   },
 });
@@ -51,7 +52,6 @@ const goBack = defineTabTool({
 
   handle: async (tab, params, response) => {
     await tab.page.goBack();
-    response.setIncludeSnapshot();
     response.addCode(`await page.goBack();`);
   },
 });
@@ -67,7 +67,6 @@ const goForward = defineTabTool({
   },
   handle: async (tab, params, response) => {
     await tab.page.goForward();
-    response.setIncludeSnapshot();
     response.addCode(`await page.goForward();`);
   },
 });

src/tools/snapshot.ts

@@ -19,20 +19,23 @@ import { z } from 'zod';
 import { defineTabTool, defineTool } from './tool.js';
 import * as javascript from '../javascript.js';
 import { generateLocator } from './utils.js';
+import { expectationSchema } from '../schemas/expectation.js';
 
 const snapshot = defineTool({
   capability: 'core',
   schema: {
     name: 'browser_snapshot',
     title: 'Page snapshot',
     description: 'Capture accessibility snapshot of the current page, this is better than screenshot',
-    inputSchema: z.object({}),
+    inputSchema: z.object({
+      expectation: expectationSchema
+    }),
     type: 'readOnly',
   },
 
   handle: async (context, params, response) => {
     await context.ensureTab();
-    response.setIncludeSnapshot();
+    // Snapshot will be handled by expectation settings
   },
 });
 
@@ -44,6 +47,7 @@ export const elementSchema = z.object({
 const clickSchema = elementSchema.extend({
   doubleClick: z.boolean().optional().describe('Whether to perform a double click instead of a single click'),
   button: z.enum(['left', 'right', 'middle']).optional().describe('Button to click, defaults to left'),
+  expectation: expectationSchema
 });
 
 const click = defineTabTool({
@@ -57,8 +61,6 @@ const click = defineTabTool({
   },
 
   handle: async (tab, params, response) => {
-    response.setIncludeSnapshot();
-
     const locator = await tab.refLocator(params);
     const button = params.button;
     const buttonAttr = button ? `{ button: '${button}' }` : '';
@@ -68,7 +70,6 @@ const click = defineTabTool({
     else
       response.addCode(`await page.${await generateLocator(locator)}.click(${buttonAttr});`);
 
-
     await tab.waitForCompletion(async () => {
       if (params.doubleClick)
         await locator.dblclick({ button });
@@ -89,13 +90,12 @@ const drag = defineTabTool({
       startRef: z.string().describe('Exact source element reference from the page snapshot'),
       endElement: z.string().describe('Human-readable target element description used to obtain the permission to interact with the element'),
       endRef: z.string().describe('Exact target element reference from the page snapshot'),
+      expectation: expectationSchema
     }),
     type: 'destructive',
   },
 
   handle: async (tab, params, response) => {
-    response.setIncludeSnapshot();
-
     const [startLocator, endLocator] = await tab.refLocators([
       { ref: params.startRef, element: params.startElement },
       { ref: params.endRef, element: params.endElement },
@@ -115,13 +115,13 @@ const hover = defineTabTool({
     name: 'browser_hover',
     title: 'Hover mouse',
     description: 'Hover over element on page',
-    inputSchema: elementSchema,
+    inputSchema: elementSchema.extend({
+      expectation: expectationSchema
+    }),
     type: 'readOnly',
   },
 
   handle: async (tab, params, response) => {
-    response.setIncludeSnapshot();
-
     const locator = await tab.refLocator(params);
     response.addCode(`await page.${await generateLocator(locator)}.hover();`);
 
@@ -133,6 +133,7 @@ const hover = defineTabTool({
 
 const selectOptionSchema = elementSchema.extend({
   values: z.array(z.string()).describe('Array of values to select in the dropdown. This can be a single value or multiple values.'),
+  expectation: expectationSchema
 });
 
 const selectOption = defineTabTool({
@@ -146,8 +147,6 @@ const selectOption = defineTabTool({
   },
 
   handle: async (tab, params, response) => {
-    response.setIncludeSnapshot();
-
     const locator = await tab.refLocator(params);
     response.addCode(`await page.${await generateLocator(locator)}.selectOption(${javascript.formatObject(params.values)});`);
 

tests/expectation-unit.spec.ts

@@ -0,0 +1,75 @@
+/**
+ * Copyright (c) Microsoft Corporation.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import { test, expect } from './fixtures.js';
+import { mergeExpectations } from '../src/schemas/expectation.js';
+
+test.describe('Expectation Unit Tests', () => {
+  test('mergeExpectations should respect user false values', () => {
+    const userExpectation = {
+      includeSnapshot: false,
+      includeConsole: false,
+      includeDownloads: false,
+      includeTabs: false,
+      includeCode: true
+    };
+    
+    const merged = mergeExpectations('browser_navigate', userExpectation);
+    
+    expect(merged.includeSnapshot).toBe(false);
+    expect(merged.includeConsole).toBe(false);
+    expect(merged.includeDownloads).toBe(false);
+    expect(merged.includeTabs).toBe(false);
+    expect(merged.includeCode).toBe(true);
+  });
+
+  test('mergeExpectations should use tool defaults when user expectation is undefined', () => {
+    const merged = mergeExpectations('browser_navigate', undefined);
+    
+    // Navigate tool defaults should be all true
+    expect(merged.includeSnapshot).toBe(true);
+    expect(merged.includeConsole).toBe(true);
+    expect(merged.includeDownloads).toBe(true);
+    expect(merged.includeTabs).toBe(true);
+    expect(merged.includeCode).toBe(true);
+  });
+
+  test('mergeExpectations should use tool defaults for click tool', () => {
+    const merged = mergeExpectations('browser_click', undefined);
+    
+    // Click tool defaults
+    expect(merged.includeSnapshot).toBe(true);
+    expect(merged.includeConsole).toBe(false);
+    expect(merged.includeDownloads).toBe(false);
+    expect(merged.includeTabs).toBe(false);
+    expect(merged.includeCode).toBe(true);
+  });
+
+  test('mergeExpectations should partially override tool defaults', () => {
+    const userExpectation = {
+      includeSnapshot: false
+      // Other values should use tool defaults
+    };
+    
+    const merged = mergeExpectations('browser_click', userExpectation);
+    
+    expect(merged.includeSnapshot).toBe(false); // User override
+    expect(merged.includeConsole).toBe(false); // Tool default
+    expect(merged.includeDownloads).toBe(false); // Tool default
+    expect(merged.includeTabs).toBe(false); // Tool default
+    expect(merged.includeCode).toBe(true); // Tool default
+  });
+});