registerTool: accept ZodType<object> for input and output schema (#816)

src/server/mcp.test.ts

@@ -3583,7 +3583,7 @@ describe('Tool title precedence', () => {
                 description: 'Tool with regular title'
             },
             async () => ({
-                content: [{ type: 'text', text: 'Response' }]
+                content: [{ type: 'text' as const, text: 'Response' }]
             })
         );
 
@@ -3598,7 +3598,7 @@ describe('Tool title precedence', () => {
                 }
             },
             async () => ({
-                content: [{ type: 'text', text: 'Response' }]
+                content: [{ type: 'text' as const, text: 'Response' }]
             })
         );
 
@@ -4162,3 +4162,239 @@ describe('elicitInput()', () => {
         ]);
     });
 });
+
+describe('Tools with union and intersection schemas', () => {
+    test('should support union schemas', async () => {
+        const server = new McpServer({
+            name: 'test',
+            version: '1.0.0'
+        });
+
+        const client = new Client({
+            name: 'test-client',
+            version: '1.0.0'
+        });
+
+        const unionSchema = z.union([
+            z.object({ type: z.literal('email'), email: z.string().email() }),
+            z.object({ type: z.literal('phone'), phone: z.string() })
+        ]);
+
+        server.registerTool('contact', { inputSchema: unionSchema }, async args => {
+            if (args.type === 'email') {
+                return {
+                    content: [{ type: 'text', text: `Email contact: ${args.email}` }]
+                };
+            } else {
+                return {
+                    content: [{ type: 'text', text: `Phone contact: ${args.phone}` }]
+                };
+            }
+        });
+
+        const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair();
+        await server.connect(serverTransport);
+        await client.connect(clientTransport);
+
+        const emailResult = await client.callTool({
+            name: 'contact',
+            arguments: {
+                type: 'email',
+                email: 'test@example.com'
+            }
+        });
+
+        expect(emailResult.content).toEqual([
+            {
+                type: 'text',
+                text: 'Email contact: test@example.com'
+            }
+        ]);
+
+        const phoneResult = await client.callTool({
+            name: 'contact',
+            arguments: {
+                type: 'phone',
+                phone: '+1234567890'
+            }
+        });
+
+        expect(phoneResult.content).toEqual([
+            {
+                type: 'text',
+                text: 'Phone contact: +1234567890'
+            }
+        ]);
+    });
+
+    test('should support intersection schemas', async () => {
+        const server = new McpServer({
+            name: 'test',
+            version: '1.0.0'
+        });
+
+        const client = new Client({
+            name: 'test-client',
+            version: '1.0.0'
+        });
+
+        const baseSchema = z.object({ id: z.string() });
+        const extendedSchema = z.object({ name: z.string(), age: z.number() });
+        const intersectionSchema = z.intersection(baseSchema, extendedSchema);
+
+        server.registerTool('user', { inputSchema: intersectionSchema }, async args => {
+            return {
+                content: [
+                    {
+                        type: 'text',
+                        text: `User: ${args.id}, ${args.name}, ${args.age} years old`
+                    }
+                ]
+            };
+        });
+
+        const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair();
+        await server.connect(serverTransport);
+        await client.connect(clientTransport);
+
+        const result = await client.callTool({
+            name: 'user',
+            arguments: {
+                id: '123',
+                name: 'John Doe',
+                age: 30
+            }
+        });
+
+        expect(result.content).toEqual([
+            {
+                type: 'text',
+                text: 'User: 123, John Doe, 30 years old'
+            }
+        ]);
+    });
+
+    test('should support complex nested schemas', async () => {
+        const server = new McpServer({
+            name: 'test',
+            version: '1.0.0'
+        });
+
+        const client = new Client({
+            name: 'test-client',
+            version: '1.0.0'
+        });
+
+        const schema = z.object({
+            items: z.array(
+                z.union([
+                    z.object({ type: z.literal('text'), content: z.string() }),
+                    z.object({ type: z.literal('number'), value: z.number() })
+                ])
+            )
+        });
+
+        server.registerTool('process', { inputSchema: schema }, async args => {
+            const processed = args.items.map(item => {
+                if (item.type === 'text') {
+                    return item.content.toUpperCase();
+                } else {
+                    return item.value * 2;
+                }
+            });
+            return {
+                content: [
+                    {
+                        type: 'text',
+                        text: `Processed: ${processed.join(', ')}`
+                    }
+                ]
+            };
+        });
+
+        const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair();
+        await server.connect(serverTransport);
+        await client.connect(clientTransport);
+
+        const result = await client.callTool({
+            name: 'process',
+            arguments: {
+                items: [
+                    { type: 'text', content: 'hello' },
+                    { type: 'number', value: 5 },
+                    { type: 'text', content: 'world' }
+                ]
+            }
+        });
+
+        expect(result.content).toEqual([
+            {
+                type: 'text',
+                text: 'Processed: HELLO, 10, WORLD'
+            }
+        ]);
+    });
+
+    test('should validate union schema inputs correctly', async () => {
+        const server = new McpServer({
+            name: 'test',
+            version: '1.0.0'
+        });
+
+        const client = new Client({
+            name: 'test-client',
+            version: '1.0.0'
+        });
+
+        const unionSchema = z.union([
+            z.object({ type: z.literal('a'), value: z.string() }),
+            z.object({ type: z.literal('b'), value: z.number() })
+        ]);
+
+        server.registerTool('union-test', { inputSchema: unionSchema }, async () => {
+            return {
+                content: [{ type: 'text', text: 'Success' }]
+            };
+        });
+
+        const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair();
+        await server.connect(serverTransport);
+        await client.connect(clientTransport);
+
+        const invalidTypeResult = await client.callTool({
+            name: 'union-test',
+            arguments: {
+                type: 'a',
+                value: 123
+            }
+        });
+
+        expect(invalidTypeResult.isError).toBe(true);
+        expect(invalidTypeResult.content).toEqual(
+            expect.arrayContaining([
+                expect.objectContaining({
+                    type: 'text',
+                    text: expect.stringContaining('Input validation error')
+                })
+            ])
+        );
+
+        const invalidDiscriminatorResult = await client.callTool({
+            name: 'union-test',
+            arguments: {
+                type: 'c',
+                value: 'test'
+            }
+        });
+
+        expect(invalidDiscriminatorResult.isError).toBe(true);
+        expect(invalidDiscriminatorResult.content).toEqual(
+            expect.arrayContaining([
+                expect.objectContaining({
+                    type: 'text',
+                    text: expect.stringContaining('Input validation error')
+                })
+            ])
+        );
+    });
+});

src/server/mcp.ts

@@ -1,6 +1,6 @@
 import { Server, ServerOptions } from './index.js';
 import { zodToJsonSchema } from 'zod-to-json-schema';
-import { z, ZodRawShape, ZodObject, ZodString, AnyZodObject, ZodTypeAny, ZodType, ZodTypeDef, ZodOptional } from 'zod';
+import { z, ZodRawShape, ZodObject, ZodString, ZodTypeAny, ZodType, ZodTypeDef, ZodOptional } from 'zod';
 import {
     Implementation,
     Tool,
@@ -650,8 +650,8 @@ export class McpServer {
         name: string,
         title: string | undefined,
         description: string | undefined,
-        inputSchema: ZodRawShape | undefined,
-        outputSchema: ZodRawShape | undefined,
+        inputSchema: ZodRawShape | ZodType<object> | undefined,
+        outputSchema: ZodRawShape | ZodType<object> | undefined,
         annotations: ToolAnnotations | undefined,
         _meta: Record<string, unknown> | undefined,
         callback: ToolCallback<ZodRawShape | undefined>
@@ -659,8 +659,8 @@ export class McpServer {
         const registeredTool: RegisteredTool = {
             title,
             description,
-            inputSchema: inputSchema === undefined ? undefined : z.object(inputSchema),
-            outputSchema: outputSchema === undefined ? undefined : z.object(outputSchema),
+            inputSchema: getZodSchemaObject(inputSchema),
+            outputSchema: getZodSchemaObject(outputSchema),
             annotations,
             _meta,
             callback,
@@ -798,7 +798,7 @@ export class McpServer {
     /**
      * Registers a tool with a config object and callback.
      */
-    registerTool<InputArgs extends ZodRawShape, OutputArgs extends ZodRawShape>(
+    registerTool<InputArgs extends ZodRawShape | ZodType<object>, OutputArgs extends ZodRawShape | ZodType<object>>(
         name: string,
         config: {
             title?: string;
@@ -1027,18 +1027,20 @@ export class ResourceTemplate {
  * - `content` if the tool does not have an outputSchema
  * - Both fields are optional but typically one should be provided
  */
-export type ToolCallback<Args extends undefined | ZodRawShape = undefined> = Args extends ZodRawShape
+export type ToolCallback<Args extends undefined | ZodRawShape | ZodType<object> = undefined> = Args extends ZodRawShape
     ? (
           args: z.objectOutputType<Args, ZodTypeAny>,
           extra: RequestHandlerExtra<ServerRequest, ServerNotification>
       ) => CallToolResult | Promise<CallToolResult>
-    : (extra: RequestHandlerExtra<ServerRequest, ServerNotification>) => CallToolResult | Promise<CallToolResult>;
+    : Args extends ZodType<infer T>
+      ? (args: T, extra: RequestHandlerExtra<ServerRequest, ServerNotification>) => CallToolResult | Promise<CallToolResult>
+      : (extra: RequestHandlerExtra<ServerRequest, ServerNotification>) => CallToolResult | Promise<CallToolResult>;
 
 export type RegisteredTool = {
     title?: string;
     description?: string;
-    inputSchema?: AnyZodObject;
-    outputSchema?: AnyZodObject;
+    inputSchema?: ZodType<object>;
+    outputSchema?: ZodType<object>;
     annotations?: ToolAnnotations;
     _meta?: Record<string, unknown>;
     callback: ToolCallback<undefined | ZodRawShape>;
@@ -1086,6 +1088,22 @@ function isZodTypeLike(value: unknown): value is ZodType {
     );
 }
 
+/**
+ * Converts a provided Zod schema to a Zod object if it is a ZodRawShape,
+ * otherwise returns the schema as is.
+ */
+function getZodSchemaObject(schema: ZodRawShape | ZodType<object> | undefined): ZodType<object> | undefined {
+    if (!schema) {
+        return undefined;
+    }
+
+    if (isZodRawShape(schema)) {
+        return z.object(schema);
+    }
+
+    return schema;
+}
+
 /**
  * Additional, optional information for annotating a resource.
  */

src/spec.types.test.ts

@@ -63,7 +63,7 @@ type MakeUnknownsNotOptional<T> =
           : T;
 
 const sdkTypeChecks = {
-    CancelledNotification: (sdk: WithJSONRPC<SDKTypes.CancelledNotification>, spec: SpecTypes.CancelledNotification) => {
+    CancelledNotification: (sdk: RemovePassthrough<WithJSONRPC<SDKTypes.CancelledNotification>>, spec: SpecTypes.CancelledNotification) => {
         sdk = spec;
         spec = sdk;
     },
@@ -75,19 +75,19 @@ const sdkTypeChecks = {
         sdk = spec;
         spec = sdk;
     },
-    ProgressNotification: (sdk: WithJSONRPC<SDKTypes.ProgressNotification>, spec: SpecTypes.ProgressNotification) => {
+    ProgressNotification: (sdk: RemovePassthrough<WithJSONRPC<SDKTypes.ProgressNotification>>, spec: SpecTypes.ProgressNotification) => {
         sdk = spec;
         spec = sdk;
     },
-    SubscribeRequest: (sdk: WithJSONRPCRequest<SDKTypes.SubscribeRequest>, spec: SpecTypes.SubscribeRequest) => {
+    SubscribeRequest: (sdk: RemovePassthrough<WithJSONRPCRequest<SDKTypes.SubscribeRequest>>, spec: SpecTypes.SubscribeRequest) => {
         sdk = spec;
         spec = sdk;
     },
-    UnsubscribeRequest: (sdk: WithJSONRPCRequest<SDKTypes.UnsubscribeRequest>, spec: SpecTypes.UnsubscribeRequest) => {
+    UnsubscribeRequest: (sdk: RemovePassthrough<WithJSONRPCRequest<SDKTypes.UnsubscribeRequest>>, spec: SpecTypes.UnsubscribeRequest) => {
         sdk = spec;
         spec = sdk;
     },
-    PaginatedRequest: (sdk: WithJSONRPCRequest<SDKTypes.PaginatedRequest>, spec: SpecTypes.PaginatedRequest) => {
+    PaginatedRequest: (sdk: RemovePassthrough<WithJSONRPCRequest<SDKTypes.PaginatedRequest>>, spec: SpecTypes.PaginatedRequest) => {
         sdk = spec;
         spec = sdk;
     },
@@ -95,7 +95,7 @@ const sdkTypeChecks = {
         sdk = spec;
         spec = sdk;
     },
-    ListRootsRequest: (sdk: WithJSONRPCRequest<SDKTypes.ListRootsRequest>, spec: SpecTypes.ListRootsRequest) => {
+    ListRootsRequest: (sdk: RemovePassthrough<WithJSONRPCRequest<SDKTypes.ListRootsRequest>>, spec: SpecTypes.ListRootsRequest) => {
         sdk = spec;
         spec = sdk;
     },
@@ -167,7 +167,7 @@ const sdkTypeChecks = {
         sdk = spec;
         spec = sdk;
     },
-    ClientNotification: (sdk: WithJSONRPC<SDKTypes.ClientNotification>, spec: SpecTypes.ClientNotification) => {
+    ClientNotification: (sdk: RemovePassthrough<WithJSONRPC<SDKTypes.ClientNotification>>, spec: SpecTypes.ClientNotification) => {
         sdk = spec;
         spec = sdk;
     },
@@ -191,7 +191,7 @@ const sdkTypeChecks = {
         sdk = spec;
         spec = sdk;
     },
-    ListToolsRequest: (sdk: WithJSONRPCRequest<SDKTypes.ListToolsRequest>, spec: SpecTypes.ListToolsRequest) => {
+    ListToolsRequest: (sdk: RemovePassthrough<WithJSONRPCRequest<SDKTypes.ListToolsRequest>>, spec: SpecTypes.ListToolsRequest) => {
         sdk = spec;
         spec = sdk;
     },
@@ -203,36 +203,42 @@ const sdkTypeChecks = {
         sdk = spec;
         spec = sdk;
     },
-    CallToolRequest: (sdk: WithJSONRPCRequest<SDKTypes.CallToolRequest>, spec: SpecTypes.CallToolRequest) => {
+    CallToolRequest: (sdk: RemovePassthrough<WithJSONRPCRequest<SDKTypes.CallToolRequest>>, spec: SpecTypes.CallToolRequest) => {
         sdk = spec;
         spec = sdk;
     },
-    ToolListChangedNotification: (sdk: WithJSONRPC<SDKTypes.ToolListChangedNotification>, spec: SpecTypes.ToolListChangedNotification) => {
+    ToolListChangedNotification: (
+        sdk: RemovePassthrough<WithJSONRPC<SDKTypes.ToolListChangedNotification>>,
+        spec: SpecTypes.ToolListChangedNotification
+    ) => {
         sdk = spec;
         spec = sdk;
     },
     ResourceListChangedNotification: (
-        sdk: WithJSONRPC<SDKTypes.ResourceListChangedNotification>,
+        sdk: RemovePassthrough<WithJSONRPC<SDKTypes.ResourceListChangedNotification>>,
         spec: SpecTypes.ResourceListChangedNotification
     ) => {
         sdk = spec;
         spec = sdk;
     },
     PromptListChangedNotification: (
-        sdk: WithJSONRPC<SDKTypes.PromptListChangedNotification>,
+        sdk: RemovePassthrough<WithJSONRPC<SDKTypes.PromptListChangedNotification>>,
         spec: SpecTypes.PromptListChangedNotification
     ) => {
         sdk = spec;
         spec = sdk;
     },
     RootsListChangedNotification: (
-        sdk: WithJSONRPC<SDKTypes.RootsListChangedNotification>,
+        sdk: RemovePassthrough<WithJSONRPC<SDKTypes.RootsListChangedNotification>>,
         spec: SpecTypes.RootsListChangedNotification
     ) => {
         sdk = spec;
         spec = sdk;
     },
-    ResourceUpdatedNotification: (sdk: WithJSONRPC<SDKTypes.ResourceUpdatedNotification>, spec: SpecTypes.ResourceUpdatedNotification) => {
+    ResourceUpdatedNotification: (
+        sdk: RemovePassthrough<WithJSONRPC<SDKTypes.ResourceUpdatedNotification>>,
+        spec: SpecTypes.ResourceUpdatedNotification
+    ) => {
         sdk = spec;
         spec = sdk;
     },
@@ -244,19 +250,25 @@ const sdkTypeChecks = {
         sdk = spec;
         spec = sdk;
     },
-    SetLevelRequest: (sdk: WithJSONRPCRequest<SDKTypes.SetLevelRequest>, spec: SpecTypes.SetLevelRequest) => {
+    SetLevelRequest: (sdk: RemovePassthrough<WithJSONRPCRequest<SDKTypes.SetLevelRequest>>, spec: SpecTypes.SetLevelRequest) => {
         sdk = spec;
         spec = sdk;
     },
-    PingRequest: (sdk: WithJSONRPCRequest<SDKTypes.PingRequest>, spec: SpecTypes.PingRequest) => {
+    PingRequest: (sdk: RemovePassthrough<WithJSONRPCRequest<SDKTypes.PingRequest>>, spec: SpecTypes.PingRequest) => {
         sdk = spec;
         spec = sdk;
     },
-    InitializedNotification: (sdk: WithJSONRPC<SDKTypes.InitializedNotification>, spec: SpecTypes.InitializedNotification) => {
+    InitializedNotification: (
+        sdk: RemovePassthrough<WithJSONRPC<SDKTypes.InitializedNotification>>,
+        spec: SpecTypes.InitializedNotification
+    ) => {
         sdk = spec;
         spec = sdk;
     },
-    ListResourcesRequest: (sdk: WithJSONRPCRequest<SDKTypes.ListResourcesRequest>, spec: SpecTypes.ListResourcesRequest) => {
+    ListResourcesRequest: (
+        sdk: RemovePassthrough<WithJSONRPCRequest<SDKTypes.ListResourcesRequest>>,
+        spec: SpecTypes.ListResourcesRequest
+    ) => {
         sdk = spec;
         spec = sdk;
     },
@@ -265,7 +277,7 @@ const sdkTypeChecks = {
         spec = sdk;
     },
     ListResourceTemplatesRequest: (
-        sdk: WithJSONRPCRequest<SDKTypes.ListResourceTemplatesRequest>,
+        sdk: RemovePassthrough<WithJSONRPCRequest<SDKTypes.ListResourceTemplatesRequest>>,
         spec: SpecTypes.ListResourceTemplatesRequest
     ) => {
         sdk = spec;
@@ -278,7 +290,10 @@ const sdkTypeChecks = {
         sdk = spec;
         spec = sdk;
     },
-    ReadResourceRequest: (sdk: WithJSONRPCRequest<SDKTypes.ReadResourceRequest>, spec: SpecTypes.ReadResourceRequest) => {
+    ReadResourceRequest: (
+        sdk: RemovePassthrough<WithJSONRPCRequest<SDKTypes.ReadResourceRequest>>,
+        spec: SpecTypes.ReadResourceRequest
+    ) => {
         sdk = spec;
         spec = sdk;
     },
@@ -314,7 +329,7 @@ const sdkTypeChecks = {
         sdk = spec;
         spec = sdk;
     },
-    ListPromptsRequest: (sdk: WithJSONRPCRequest<SDKTypes.ListPromptsRequest>, spec: SpecTypes.ListPromptsRequest) => {
+    ListPromptsRequest: (sdk: RemovePassthrough<WithJSONRPCRequest<SDKTypes.ListPromptsRequest>>, spec: SpecTypes.ListPromptsRequest) => {
         sdk = spec;
         spec = sdk;
     },
@@ -322,7 +337,7 @@ const sdkTypeChecks = {
         sdk = spec;
         spec = sdk;
     },
-    GetPromptRequest: (sdk: WithJSONRPCRequest<SDKTypes.GetPromptRequest>, spec: SpecTypes.GetPromptRequest) => {
+    GetPromptRequest: (sdk: RemovePassthrough<WithJSONRPCRequest<SDKTypes.GetPromptRequest>>, spec: SpecTypes.GetPromptRequest) => {
         sdk = spec;
         spec = sdk;
     },
@@ -418,13 +433,16 @@ const sdkTypeChecks = {
         spec = sdk;
     },
     LoggingMessageNotification: (
-        sdk: MakeUnknownsNotOptional<WithJSONRPC<SDKTypes.LoggingMessageNotification>>,
+        sdk: RemovePassthrough<MakeUnknownsNotOptional<WithJSONRPC<SDKTypes.LoggingMessageNotification>>>,
         spec: SpecTypes.LoggingMessageNotification
     ) => {
         sdk = spec;
         spec = sdk;
     },
-    ServerNotification: (sdk: MakeUnknownsNotOptional<WithJSONRPC<SDKTypes.ServerNotification>>, spec: SpecTypes.ServerNotification) => {
+    ServerNotification: (
+        sdk: RemovePassthrough<MakeUnknownsNotOptional<WithJSONRPC<SDKTypes.ServerNotification>>>,
+        spec: SpecTypes.ServerNotification
+    ) => {
         sdk = spec;
         spec = sdk;
     },
@@ -451,6 +469,26 @@ const MISSING_SDK_TYPES = [
     'Role',
     'Error', // The inner error object of a JSONRPCError
 
+    // Params types are not exported as standalone types in the SDK (they're inferred from Zod schemas):
+    'RequestParams',
+    'NotificationParams',
+    'CancelledNotificationParams',
+    'InitializeRequestParams',
+    'ProgressNotificationParams',
+    'PaginatedRequestParams',
+    'ResourceRequestParams',
+    'ReadResourceRequestParams',
+    'SubscribeRequestParams',
+    'UnsubscribeRequestParams',
+    'SetLevelRequestParams',
+    'GetPromptRequestParams',
+    'CompleteRequestParams',
+    'CallToolRequestParams',
+    'CreateMessageRequestParams',
+    'LoggingMessageNotificationParams',
+    'ResourceUpdatedNotificationParams',
+    'ElicitRequestParams',
+
     // These aren't supported by the SDK yet:
     // TODO: Add definitions to the SDK
     'Annotations',
@@ -470,7 +508,7 @@ describe('Spec Types', () => {
     it('should define some expected types', () => {
         expect(specTypes).toContain('JSONRPCNotification');
         expect(specTypes).toContain('ElicitResult');
-        expect(specTypes).toHaveLength(94);
+        expect(specTypes).toHaveLength(112);
     });
 
     it('should have up to date list of missing sdk types', () => {
@@ -479,10 +517,16 @@ describe('Spec Types', () => {
         }
     });
 
-    describe('Compatibility tests', () => {
-        it.each(typesToCheck)('%s should have a compatibility test', type => {
-            expect(sdkTypeChecks[type as keyof typeof sdkTypeChecks]).toBeDefined();
-        });
+    it('should have comprehensive compatibility tests', () => {
+        const missingTests = [];
+
+        for (const typeName of typesToCheck) {
+            if (!sdkTypeChecks[typeName as keyof typeof sdkTypeChecks]) {
+                missingTests.push(typeName);
+            }
+        }
+
+        expect(missingTests).toHaveLength(0);
     });
 
     describe('Missing SDK Types', () => {