refactor: prepare for SDK v2 (#2771)

The code changes here are a subset that works for both the current and
the update one. Splitting to reduce the review process.

scripts/eval_gemini.ts

@@ -10,9 +10,8 @@ import {pathToFileURL} from 'node:url';
 import {parseArgs} from 'node:util';
 
 import {GoogleGenAI, mcpToTool} from '@google/genai';
-import {Client} from '@modelcontextprotocol/sdk/client/index.js';
-import {StdioClientTransport} from '@modelcontextprotocol/sdk/client/stdio.js';
 
+import {Client, StdioClientTransport} from '../build/src/third_party/index.js';
 import {TestServer} from '../build/tests/server.js';
 
 const ROOT_DIR = path.resolve(import.meta.dirname, '..');

scripts/generate-cli.ts

@@ -7,9 +7,7 @@
 import fs from 'node:fs';
 import path from 'node:path';
 
-import {Client} from '@modelcontextprotocol/sdk/client/index.js';
-import {StdioClientTransport} from '@modelcontextprotocol/sdk/client/stdio.js';
-
+import {Client, StdioClientTransport} from '../build/src/third_party/index.js';
 import {mcpOptions, parseArguments} from '../build/src/config/mcp-options.js';
 import {
   isCategoryOffByDefault,

scripts/generate-docs.ts

@@ -6,7 +6,6 @@
 
 import fs from 'node:fs';
 
-import type {Tool} from '@modelcontextprotocol/sdk/types.js';
 import type {Options as YargsOptions} from 'yargs';
 
 import {
@@ -17,7 +16,12 @@ import {
   isCategoryOffByDefault,
   categoryToFlagName,
 } from '../build/src/config/category-options.js';
+import type {Tool, zod} from '../src/third_party/index.js';
 import {ToolCategory, labels} from '../build/src/tools/categories.js';
+import type {
+  DefinedPageTool,
+  ToolDefinition,
+} from '../src/tools/ToolDefinition.js';
 import {createTools} from '../build/src/tools/tools.js';
 
 const OUTPUT_PATH = './docs/tool-reference.md';
@@ -37,18 +41,21 @@ interface ZodCheck {
 }
 
 interface ZodDef {
-  typeName: string;
+  typeName?: string;
   checks?: ZodCheck[];
   values?: string[];
-  type?: ZodSchema;
+  entries?: Record<string, string>;
+  type?: string | ZodSchema;
   innerType?: ZodSchema;
   schema?: ZodSchema;
-  defaultValue?: () => unknown;
+  in?: ZodSchema;
+  defaultValue?: (() => unknown) | unknown;
 }
 
 interface ZodSchema {
   _def: ZodDef;
   description?: string;
+  isInt?: boolean;
 }
 
 interface TypeInfo {
@@ -200,21 +207,35 @@ function getZodTypeInfo(schema: ZodSchema): TypeInfo {
   let def = schema._def;
   let defaultValue: unknown;
 
+  let typeName = def.typeName;
+  if (!typeName && typeof def.type === 'string') {
+    typeName = 'Zod' + def.type.charAt(0).toUpperCase() + def.type.slice(1);
+  }
+
   // Unwrap optional/default/effects
   while (
-    def.typeName === 'ZodOptional' ||
-    def.typeName === 'ZodDefault' ||
-    def.typeName === 'ZodEffects'
+    typeName === 'ZodOptional' ||
+    typeName === 'ZodDefault' ||
+    typeName === 'ZodEffects' ||
+    typeName === 'ZodPipeline' ||
+    typeName === 'ZodPipe'
   ) {
-    if (def.typeName === 'ZodDefault' && def.defaultValue) {
-      defaultValue = def.defaultValue();
+    if (typeName === 'ZodDefault' && def.defaultValue !== undefined) {
+      defaultValue =
+        typeof def.defaultValue === 'function'
+          ? (def.defaultValue as () => unknown)()
+          : def.defaultValue;
     }
-    const next = def.innerType || def.schema;
+    const next = def.innerType || def.schema || def.in;
     if (!next) {
       break;
     }
     schema = next;
     def = schema._def;
+    typeName = def.typeName;
+    if (!typeName && typeof def.type === 'string') {
+      typeName = 'Zod' + def.type.charAt(0).toUpperCase() + def.type.slice(1);
+    }
     if (!description && schema.description) {
       description = schema.description;
     }
@@ -228,14 +249,15 @@ function getZodTypeInfo(schema: ZodSchema): TypeInfo {
     result.default = defaultValue;
   }
 
-  switch (def.typeName) {
+  switch (typeName) {
     case 'ZodString':
       result.type = 'string';
       break;
     case 'ZodNumber':
-      result.type = def.checks?.some((c: ZodCheck) => c.kind === 'int')
-        ? 'integer'
-        : 'number';
+      result.type =
+        schema.isInt || def.checks?.some((c: ZodCheck) => c.kind === 'int')
+          ? 'integer'
+          : 'number';
       break;
     case 'ZodBoolean':
       result.type = 'boolean';
@@ -243,10 +265,13 @@ function getZodTypeInfo(schema: ZodSchema): TypeInfo {
     case 'ZodEnum':
       result.type = 'string';
       result.enum = def.values;
+      if (!result.enum && def.entries) {
+        result.enum = Object.values(def.entries);
+      }
       break;
     case 'ZodArray':
       result.type = 'array';
-      if (def.type) {
+      if (typeof def.type !== 'string' && def.type) {
         result.items = getZodTypeInfo(def.type);
       }
       break;
@@ -258,14 +283,27 @@ function getZodTypeInfo(schema: ZodSchema): TypeInfo {
 
 function isRequired(schema: ZodSchema): boolean {
   let def = schema._def;
-  while (def.typeName === 'ZodEffects') {
-    if (!def.schema) {
+  let typeName = def.typeName;
+  if (!typeName && typeof def.type === 'string') {
+    typeName = 'Zod' + def.type.charAt(0).toUpperCase() + def.type.slice(1);
+  }
+  while (
+    typeName === 'ZodEffects' ||
+    typeName === 'ZodPipeline' ||
+    typeName === 'ZodPipe'
+  ) {
+    const next = def.schema || def.in;
+    if (!next) {
       break;
     }
-    schema = def.schema;
+    schema = next;
     def = schema._def;
+    typeName = def.typeName;
+    if (!typeName && typeof def.type === 'string') {
+      typeName = 'Zod' + def.type.charAt(0).toUpperCase() + def.type.slice(1);
+    }
   }
-  return def.typeName !== 'ZodOptional' && def.typeName !== 'ZodDefault';
+  return typeName !== 'ZodOptional' && typeName !== 'ZodDefault';
 }
 
 async function generateReference(
@@ -412,8 +450,9 @@ async function generateReference(
   );
 }
 
-// eslint-disable-next-line @typescript-eslint/no-explicit-any
-function getToolsAndCategories(tools: any) {
+function getToolsAndCategories(
+  tools: Array<ToolDefinition<zod.ZodRawShape> | DefinedPageTool>,
+) {
   // Convert ToolDefinitions to ToolWithAnnotations
   const toolsWithAnnotations: ToolWithAnnotations[] = tools
     .filter(tool => {

src/SlimMcpResponse.ts

@@ -4,13 +4,9 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import type {
-  TextContent,
-  ImageContent,
-} from '@modelcontextprotocol/sdk/types.js';
-
 import type {McpContext} from './McpContext.js';
 import {McpResponse} from './McpResponse.js';
+import type {TextContent, ImageContent} from './third_party/index.js';
 
 export class SlimMcpResponse extends McpResponse {
   override async handle(_context: McpContext): Promise<{

src/telemetry/transformation.ts

@@ -53,13 +53,36 @@ const SUPPORTED_ZOD_TYPES = [
 ] as const;
 type ZodType = (typeof SUPPORTED_ZOD_TYPES)[number];
 
+interface ZodLikeSchema {
+  _def: {
+    typeName?: string;
+    type?: string;
+    innerType?: unknown;
+    schema?: unknown;
+    in?: unknown;
+    out?: unknown;
+    values?: unknown[];
+    entries?: Record<string, unknown>;
+  };
+}
+
+function isObjectWithDef(val: unknown): val is ZodLikeSchema {
+  return typeof val === 'object' && val !== null && '_def' in val;
+}
+
 function isZodType(type: string): type is ZodType {
   return SUPPORTED_ZOD_TYPES.includes(type as ZodType);
 }
 
-export function getZodType(zodType: zod.ZodTypeAny): ZodType {
+export function getZodType(zodType: unknown): ZodType {
+  if (!isObjectWithDef(zodType)) {
+    throw new Error('Invalid zod schema');
+  }
   const def = zodType._def;
-  const typeName = def.typeName;
+  let typeName = def.typeName;
+  if (!typeName && def.type) {
+    typeName = 'Zod' + def.type.charAt(0).toUpperCase() + def.type.slice(1);
+  }
 
   if (
     typeName === 'ZodOptional' ||
@@ -71,8 +94,17 @@ export function getZodType(zodType: zod.ZodTypeAny): ZodType {
   if (typeName === 'ZodEffects') {
     return getZodType(def.schema);
   }
+  if (typeName === 'ZodPipeline' || typeName === 'ZodPipe') {
+    return getZodType(
+      isObjectWithDef(def.in) &&
+        (def.in._def.type === 'transform' ||
+          def.in._def.typeName === 'ZodTransform')
+        ? def.out
+        : def.in,
+    );
+  }
 
-  if (isZodType(typeName)) {
+  if (typeName && isZodType(typeName)) {
     return typeName;
   }
   throw new Error(`Unsupported zod type for tool parameter: ${typeName}`);
@@ -82,9 +114,15 @@ export function getZodType(zodType: zod.ZodTypeAny): ZodType {
  * Resolves the values of an enum parameter, unwrapping any optional/default/
  * nullable/effects wrappers (in any order), mirroring {@link getZodType}.
  */
-export function getEnumValues(zodType: zod.ZodTypeAny): unknown[] {
+export function getEnumValues(zodType: unknown): unknown[] {
+  if (!isObjectWithDef(zodType)) {
+    throw new Error('Invalid zod schema');
+  }
   const def = zodType._def;
-  const typeName = def.typeName;
+  let typeName = def.typeName;
+  if (!typeName && def.type) {
+    typeName = 'Zod' + def.type.charAt(0).toUpperCase() + def.type.slice(1);
+  }
 
   if (
     typeName === 'ZodOptional' ||
@@ -96,8 +134,22 @@ export function getEnumValues(zodType: zod.ZodTypeAny): unknown[] {
   if (typeName === 'ZodEffects') {
     return getEnumValues(def.schema);
   }
+  if (typeName === 'ZodPipeline' || typeName === 'ZodPipe') {
+    return getEnumValues(
+      isObjectWithDef(def.in) &&
+        (def.in._def.type === 'transform' ||
+          def.in._def.typeName === 'ZodTransform')
+        ? def.out
+        : def.in,
+    );
+  }
   if (typeName === 'ZodEnum') {
-    return def.values;
+    if (def.values) {
+      return def.values;
+    }
+    if (def.entries) {
+      return Object.values(def.entries);
+    }
   }
   throw new Error(`Cannot resolve enum values for zod type: ${typeName}`);
 }

src/third_party/index.ts

@@ -10,6 +10,7 @@ import 'core-js/proposals/iterator-helpers.js';
 
 import type {Flags, OutputMode, Result, RunnerResult} from 'lighthouse';
 import type {Page} from 'puppeteer-core';
+import {z as zod} from 'zod';
 
 export type {Flags, Result, RunnerResult, OutputMode};
 
@@ -18,22 +19,26 @@ export {default as yargs} from 'yargs';
 export {hideBin} from 'yargs/helpers';
 export {default as semver} from 'semver';
 export {McpServer} from '@modelcontextprotocol/sdk/server/mcp.js';
-export {type ShapeOutput} from '@modelcontextprotocol/sdk/server/zod-compat.js';
 export {StdioServerTransport} from '@modelcontextprotocol/sdk/server/stdio.js';
 export {StdioClientTransport} from '@modelcontextprotocol/sdk/client/stdio.js';
 export {Client} from '@modelcontextprotocol/sdk/client/index.js';
 export type {Transport} from '@modelcontextprotocol/sdk/shared/transport.js';
 export {
   type CallToolResult,
+  type ClientCapabilities,
   SetLevelRequestSchema,
   type ImageContent,
   type TextContent,
+  type Tool,
   type Root,
   ListRootsRequestSchema,
   RootsListChangedNotificationSchema,
   ListRootsResultSchema,
 } from '@modelcontextprotocol/sdk/types.js';
-export {z as zod} from 'zod';
+export {zod};
+export type ShapeOutput<T extends zod.ZodRawShape> = zod.output<
+  zod.ZodObject<T>
+>;
 export {default as ajv} from 'ajv';
 export {
   Locator,

src/tools/ToolDefinition.ts

@@ -56,14 +56,18 @@ export type FileVerificationOption =
 
 type AllKeys<T> = T extends unknown ? keyof T : never;
 
+type ExtractSchemaField<Schema, K extends PropertyKey> = Schema extends unknown
+  ? K extends keyof Schema
+    ? Exclude<Schema[K], undefined>
+    : never
+  : never;
+
 export type MergeSchema<Schema extends zod.ZodRawShape> = {
-  [K in AllKeys<Schema>]: Schema extends unknown
-    ? K extends keyof Schema
-      ? undefined extends Schema[K]
-        ? Exclude<Schema[K], undefined> | zod.ZodUndefined
-        : Schema[K]
-      : zod.ZodUndefined
-    : never;
+  [K in AllKeys<Schema>]: K extends keyof Schema
+    ? undefined extends Schema[K]
+      ? zod.ZodOptional<ExtractSchemaField<Schema, K>>
+      : Schema[K]
+    : zod.ZodOptional<ExtractSchemaField<Schema, K>>;
 };
 
 export interface BaseToolDefinition<
@@ -98,8 +102,12 @@ export interface ToolDefinition<
   ): Promise<void>;
 }
 
+export type SchemaType<T extends zod.ZodRawShape> = zod.output<
+  zod.ZodObject<MergeSchema<T>>
+>;
+
 export interface Request<Schema extends zod.ZodRawShape> {
-  params: zod.objectOutputType<MergeSchema<Schema>, zod.ZodTypeAny>;
+  params: SchemaType<Schema>;
 }
 
 export interface ImageContentData {
@@ -454,13 +462,13 @@ export const timeoutSchema = {
   timeout: zod
     .number()
     .int()
+    .transform(value => {
+      return value <= 0 ? undefined : value;
+    })
     .optional()
     .describe(
       `Maximum wait time in milliseconds. If set to 0, the default timeout will be used.`,
-    )
-    .transform(value => {
-      return value && value <= 0 ? undefined : value;
-    }),
+    ),
 };
 
 export function viewportTransform(arg: string | undefined):

src/tools/emulation.ts

@@ -67,8 +67,8 @@ export const emulate = definePageTool(() => ({
       ),
     geolocation: zod
       .string()
-      .optional()
       .transform(geolocationTransform)
+      .optional()
       .describe(
         'Geolocation (`<latitude>,<longitude>`) to emulate. Latitude between -90 and 90. Longitude between -180 and 180. Omit to clear the geolocation override.',
       ),
@@ -86,15 +86,15 @@ export const emulate = definePageTool(() => ({
       ),
     viewport: zod
       .string()
-      .optional()
       .transform(viewportTransform)
+      .optional()
       .describe(
         `Emulate device viewports '<width>x<height>x<devicePixelRatio>[,mobile][,touch][,landscape]'. 'touch' and 'mobile' to emulate mobile devices. 'landscape' to emulate landscape mode.`,
       ),
     extraHttpHeaders: zod
       .string()
-      .optional()
       .transform(headerStringTransform)
+      .optional()
       .describe(
         'Extra HTTP headers as a JSON string object, e.g. {"X-Custom": "value", "Authorization": "Bearer token"}. Headers are included into every HTTP request originating from the page and persist across navigations until cleared. Pass an empty string to clear all extra headers.',
       ),

src/tools/input.ts

@@ -394,15 +394,18 @@ export const fillForm = definePageTool(() => ({
   schema: {
     elements: zod
       .array(
-        // eslint-disable-next-line @local/enforce-zod-schema
-        zod.object({
-          uid: zod.string().describe('The uid of the element to fill out'),
-          value: zod
-            .string()
-            .describe(
-              'Value for the element. "true" or "false" for checkboxes and toggles, "true" for radio buttons.',
-            ),
-        }),
+        /* eslint-disable @local/enforce-zod-schema */
+        zod
+          .object({
+            uid: zod.string().describe('The uid of the element to fill out'),
+            value: zod
+              .string()
+              .describe(
+                'Value for the element. "true" or "false" for checkboxes and toggles, "true" for radio buttons.',
+              ),
+          })
+          .describe('An element to fill out'),
+        /* eslint-enable @local/enforce-zod-schema */
       )
       .describe('Elements from snapshot to fill out.'),
     includeSnapshot: includeSnapshotSchema,

tests/index.test.ts

@@ -11,18 +11,18 @@ import path from 'node:path';
 import {describe, it} from 'node:test';
 import {pathToFileURL} from 'node:url';
 
-import {Client} from '@modelcontextprotocol/sdk/client/index.js';
-import {StdioClientTransport} from '@modelcontextprotocol/sdk/client/stdio.js';
-import {
-  ListRootsRequestSchema,
-  RootsListChangedNotificationSchema,
-  type ClientCapabilities,
-  type TextContent,
-} from '@modelcontextprotocol/sdk/types.js';
 import {executablePath} from 'puppeteer';
 
 import {mcpOptions} from '../src/config/mcp-options.js';
 import {getOffByDefaultCategories} from '../src/config/category-options.js';
+import {
+  Client,
+  ListRootsRequestSchema,
+  RootsListChangedNotificationSchema,
+  StdioClientTransport,
+  type ClientCapabilities,
+  type TextContent,
+} from '../src/third_party/index.js';
 import type {ToolCategory} from '../src/tools/categories.js';
 import type {ToolDefinition} from '../src/tools/ToolDefinition.js';
 
@@ -520,13 +520,11 @@ async function getToolsWithFilteredCategories(
     }
     const fileTools = await import(`../src/tools/${file}`);
 
-    for (const maybeTool of Object.values<unknown>(fileTools)) {
-      let tool;
-      if (typeof maybeTool === 'function') {
-        tool = (maybeTool as (val: boolean) => ToolDefinition)(false);
-      } else {
-        tool = maybeTool as ToolDefinition;
+    for (const maybeTool of Object.values(fileTools)) {
+      if (typeof maybeTool !== 'function') {
+        continue;
       }
+      const tool = maybeTool({});
 
       // Skipping all files that are not tool files
       if (tool === null || typeof tool !== 'object' || !('name' in tool)) {

tests/third_party_notices.test.js.snapshot

@@ -26,6 +26,36 @@ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 THE SOFTWARE.
 
 
+-------------------- DEPENDENCY DIVIDER --------------------
+
+Name: zod
+URL: https://zod.dev
+Version: <VERSION>
+License: MIT
+
+MIT License
+
+Copyright (c) 2025 Colin McDonnell
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+
+
 -------------------- DEPENDENCY DIVIDER --------------------
 
 Name: cliui
@@ -318,36 +348,6 @@ ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR
 IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 
 
--------------------- DEPENDENCY DIVIDER --------------------
-
-Name: zod
-URL: https://zod.dev
-Version: <VERSION>
-License: MIT
-
-MIT License
-
-Copyright (c) 2025 Colin McDonnell
-
-Permission is hereby granted, free of charge, to any person obtaining a copy
-of this software and associated documentation files (the "Software"), to deal
-in the Software without restriction, including without limitation the rights
-to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-copies of the Software, and to permit persons to whom the Software is
-furnished to do so, subject to the following conditions:
-
-The above copyright notice and this permission notice shall be included in all
-copies or substantial portions of the Software.
-
-THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-SOFTWARE.
-
-
 -------------------- DEPENDENCY DIVIDER --------------------
 
 Name: @modelcontextprotocol/sdk

tests/tools/slim/tools.test.ts

@@ -175,7 +175,7 @@ describe('slim', () => {
       const page = context.getSelectedMcpPage().pptrPage;
       await page.setContent(fixture.html);
       await screenshot(args).handler(
-        {params: {format: 'png'}, page: context.getSelectedMcpPage()},
+        {params: {}, page: context.getSelectedMcpPage()},
         response,
         context,
       );

tests/utils.ts

@@ -8,10 +8,10 @@ import assert from 'node:assert';
 import {spawn, type ChildProcess} from 'node:child_process';
 import path from 'node:path';
 
-import type {CallToolResult} from '@modelcontextprotocol/sdk/types.js';
-
 import type {Browser} from 'puppeteer';
 import puppeteer, {Locator} from 'puppeteer';
+
+import type {CallToolResult} from '../src/third_party/index.js';
 import type {
   Frame,
   HTTPRequest,