refactor: improve type safety for page scoped tools

Author: OrKoN

src/server.ts

@@ -23,7 +23,10 @@ import {
   SetLevelRequestSchema,
 } from './third_party/index.js';
 import {ToolCategory} from './tools/categories.js';
-import type {ToolDefinition} from './tools/ToolDefinition.js';
+import type {
+  DefinedPageTool,
+  ToolDefinition,
+} from './tools/ToolDefinition.js';
 import {pageIdSchema} from './tools/ToolDefinition.js';
 import {createTools} from './tools/tools.js';
 import {VERSION} from './version.js';
@@ -107,7 +110,7 @@ export async function createMcpServer(
 
   const toolMutex = new Mutex();
 
-  function registerTool(tool: ToolDefinition): void {
+  function registerTool(tool: ToolDefinition | DefinedPageTool): void {
     if (
       tool.annotations.category === ToolCategory.EMULATION &&
       serverArgs.categoryEmulation === false
@@ -151,7 +154,9 @@ export async function createMcpServer(
       return;
     }
     const schema =
-      tool.annotations.pageScoped && serverArgs.experimentalPageIdRouting
+      'pageScoped' in tool &&
+      tool.pageScoped &&
+      serverArgs.experimentalPageIdRouting
         ? {...tool.schema, ...pageIdSchema}
         : tool.schema;
 
@@ -174,22 +179,31 @@ export async function createMcpServer(
           const response = serverArgs.slim
             ? new SlimMcpResponse(serverArgs)
             : new McpResponse(serverArgs);
-          const page =
-            tool.annotations.pageScoped && serverArgs.experimentalPageIdRouting
-              ? context.resolvePageById(params.pageId as number | undefined)
-              : undefined;
-          if (page) {
-            context.setRequestPage(page);
-          }
           try {
-            await tool.handler(
-              {
-                params,
-                page,
-              },
-              response,
-              context,
-            );
+            if ('pageScoped' in tool && tool.pageScoped) {
+              const page =
+                serverArgs.experimentalPageIdRouting && params.pageId
+                  ? context.resolvePageById(params.pageId)
+                  : context.getSelectedPage();
+              context.setRequestPage(page);
+              await tool.handler(
+                {
+                  params,
+                  page,
+                },
+                response,
+                context,
+              );
+            } else {
+              await tool.handler(
+                // @ts-expect-error types do not match.
+                {
+                  params,
+                },
+                response,
+                context,
+              );
+            }
             const {content, structuredContent} = await response.handle(
               tool.name,
               context,

src/tools/ToolDefinition.ts

@@ -20,7 +20,7 @@ import type {PaginationOptions} from '../utils/types.js';
 
 import type {ToolCategory} from './categories.js';
 
-export interface ToolDefinition<
+export interface BaseToolDefinition<
   Schema extends zod.ZodRawShape = zod.ZodRawShape,
 > {
   name: string;
@@ -33,14 +33,14 @@ export interface ToolDefinition<
      */
     readOnlyHint: boolean;
     conditions?: string[];
-    /**
-     * If true, the tool operates on a specific page.
-     * The `pageId` schema field is auto-injected and the resolved
-     * page is provided via `request.page`.
-     */
-    pageScoped?: boolean;
   };
   schema: Schema;
+}
+
+export interface ToolDefinition<
+  Schema extends zod.ZodRawShape = zod.ZodRawShape,
+> extends BaseToolDefinition<Schema> {
+  schema: Schema;
   handler: (
     request: Request<Schema>,
     response: Response,
@@ -50,8 +50,6 @@ export interface ToolDefinition<
 
 export interface Request<Schema extends zod.ZodRawShape> {
   params: zod.objectOutputType<Schema, zod.ZodTypeAny>;
-  /** Populated centrally for tools with `pageScoped: true`. */
-  page?: Page;
 }
 
 export interface ImageContentData {
@@ -215,28 +213,66 @@ export function defineTool<
   if (typeof definition === 'function') {
     const factory = definition;
     return (args: Args) => {
-      const tool = factory(args);
-      wrapPageScopedHandler(tool);
-      return tool;
+      return factory(args);
     };
   }
-  wrapPageScopedHandler(definition);
   return definition;
 }
 
-function wrapPageScopedHandler<Schema extends zod.ZodRawShape>(
-  definition: ToolDefinition<Schema>,
-) {
-  if (definition.annotations.pageScoped) {
-    const originalHandler = definition.handler;
-    definition.handler = async (request, response, context) => {
-      // In production, main.ts resolves request.page centrally before calling
-      // the handler. This fallback exists for tests that invoke handlers
-      // directly without going through main.ts.
-      request.page ??= context.getSelectedPage();
-      return originalHandler(request, response, context);
+interface PageToolDefinition<
+  Schema extends zod.ZodRawShape = zod.ZodRawShape,
+> extends BaseToolDefinition<Schema> {
+  handler: (
+    request: Request<Schema> & {page: Page},
+    response: Response,
+    context: Context,
+  ) => Promise<void>;
+}
+
+export type DefinedPageTool<
+  Schema extends zod.ZodRawShape = zod.ZodRawShape,
+> = PageToolDefinition<Schema> & {
+  pageScoped: true;
+  handler: (
+    request: Request<Schema> & {page: Page},
+    response: Response,
+    context: Context,
+  ) => Promise<void>;
+};
+
+export function definePageTool<Schema extends zod.ZodRawShape>(
+  definition: PageToolDefinition<Schema>,
+): DefinedPageTool<Schema>;
+
+export function definePageTool<
+  Schema extends zod.ZodRawShape,
+  Args extends ParsedArguments = ParsedArguments,
+>(
+  definition: (args?: Args) => PageToolDefinition<Schema>,
+): (args?: Args) => DefinedPageTool<Schema>;
+
+export function definePageTool<
+  Schema extends zod.ZodRawShape,
+  Args extends ParsedArguments = ParsedArguments,
+>(
+  definition:
+    | PageToolDefinition<Schema>
+    | ((args?: Args) => PageToolDefinition<Schema>),
+): DefinedPageTool<Schema> | ((args?: Args) => DefinedPageTool<Schema>) {
+  if (typeof definition === 'function') {
+    return (args?: Args): DefinedPageTool<Schema> => {
+      const tool = definition(args);
+      return {
+        ...tool,
+        pageScoped: true,
+      };
     };
   }
+
+  return {
+    ...definition,
+    pageScoped: true,
+  } as DefinedPageTool<Schema>;
 }
 
 export const CLOSE_PAGE_ERROR =

src/tools/console.ts

@@ -8,7 +8,7 @@ import {zod} from '../third_party/index.js';
 import type {ConsoleMessageType} from '../third_party/index.js';
 
 import {ToolCategory} from './categories.js';
-import {defineTool} from './ToolDefinition.js';
+import {definePageTool} from './ToolDefinition.js';
 type ConsoleResponseType = ConsoleMessageType | 'issue';
 
 const FILTERABLE_MESSAGE_TYPES: [
@@ -37,14 +37,13 @@ const FILTERABLE_MESSAGE_TYPES: [
   'issue',
 ];
 
-export const listConsoleMessages = defineTool({
+export const listConsoleMessages = definePageTool({
   name: 'list_console_messages',
   description:
     'List all console messages for the currently selected page since the last navigation.',
   annotations: {
     category: ToolCategory.DEBUGGING,
     readOnlyHint: true,
-    pageScoped: true,
   },
   schema: {
     pageSize: zod
@@ -87,13 +86,12 @@ export const listConsoleMessages = defineTool({
   },
 });
 
-export const getConsoleMessage = defineTool({
+export const getConsoleMessage = definePageTool({
   name: 'get_console_message',
   description: `Gets a console message by its ID. You can get all messages by calling ${listConsoleMessages.name}.`,
   annotations: {
     category: ToolCategory.DEBUGGING,
     readOnlyHint: true,
-    pageScoped: true,
   },
   schema: {
     msgid: zod

src/tools/emulation.ts

@@ -8,21 +8,20 @@
 import {zod, PredefinedNetworkConditions} from '../third_party/index.js';
 
 import {ToolCategory} from './categories.js';
-import {defineTool} from './ToolDefinition.js';
+import {definePageTool} from './ToolDefinition.js';
 
 const throttlingOptions: [string, ...string[]] = [
   'No emulation',
   'Offline',
   ...Object.keys(PredefinedNetworkConditions),
 ];
 
-export const emulate = defineTool({
+export const emulate = definePageTool({
   name: 'emulate',
   description: `Emulates various features on the selected page.`,
   annotations: {
     category: ToolCategory.EMULATION,
     readOnlyHint: false,
-    pageScoped: true,
   },
   schema: {
     networkConditions: zod
@@ -105,7 +104,7 @@ export const emulate = defineTool({
       ),
   },
   handler: async (request, _response, context) => {
-    const page = request.page!;
+    const page = request.page;
     await context.emulate(request.params, page);
   },
 });