Skip to content

Commit 27e4ddf

Browse files
feat(compat): registerTool/registerPrompt accept raw Zod shape (auto-wrap with z.object)
1 parent 9ed62fe commit 27e4ddf

File tree

4 files changed

+163
-3
lines changed

4 files changed

+163
-3
lines changed
Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
---
2+
'@modelcontextprotocol/core': patch
3+
'@modelcontextprotocol/server': patch
4+
---
5+
6+
`registerTool`/`registerPrompt` accept a raw Zod shape (`{ field: z.string() }`) for `inputSchema`/`outputSchema`/`argsSchema` in addition to a wrapped Standard Schema. Raw shapes are auto-wrapped with `z.object()`. Both forms are first-class.

packages/core/src/util/standardSchema.ts

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,8 @@
66

77
/* eslint-disable @typescript-eslint/no-namespace */
88

9+
import * as z from 'zod/v4';
10+
911
// Standard Schema interfaces — vendored from https://standardschema.dev (spec v1, Jan 2025)
1012

1113
export interface StandardTypedV1<Input = unknown, Output = Input> {
@@ -136,6 +138,38 @@ export function isStandardSchemaWithJSON(schema: unknown): schema is StandardSch
136138
return isStandardJSONSchema(schema) && isStandardSchema(schema);
137139
}
138140

141+
/**
142+
* Detects a "raw shape" — a plain object whose values are Zod (or other
143+
* Standard Schema) field schemas, e.g. `{ name: z.string() }`. Powers the
144+
* auto-wrap in {@linkcode normalizeRawShapeSchema}.
145+
*
146+
* @internal
147+
*/
148+
export function isZodRawShape(obj: unknown): obj is Record<string, StandardSchemaV1> {
149+
if (typeof obj !== 'object' || obj === null) return false;
150+
if (isStandardSchema(obj)) return false;
151+
const values = Object.values(obj);
152+
return values.length > 0 && values.every(v => isStandardSchema(v) || (typeof v === 'object' && v !== null && '_def' in v));
153+
}
154+
155+
/**
156+
* Accepts either a {@linkcode StandardSchemaWithJSON} or a raw Zod shape
157+
* `{ field: z.string() }` and returns a {@linkcode StandardSchemaWithJSON}.
158+
* Raw shapes are wrapped with `z.object()` so the rest of the pipeline sees a
159+
* uniform schema type; already-wrapped schemas pass through unchanged.
160+
*
161+
* @internal
162+
*/
163+
export function normalizeRawShapeSchema(
164+
schema: StandardSchemaWithJSON | Record<string, StandardSchemaV1> | undefined
165+
): StandardSchemaWithJSON | undefined {
166+
if (schema === undefined) return undefined;
167+
if (isZodRawShape(schema)) {
168+
return z.object(schema as z.ZodRawShape) as StandardSchemaWithJSON;
169+
}
170+
return schema;
171+
}
172+
139173
// JSON Schema conversion
140174

141175
/**

packages/server/src/server/mcp.ts

Lines changed: 69 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@ import type {
3030
import {
3131
assertCompleteRequestPrompt,
3232
assertCompleteRequestResourceTemplate,
33+
normalizeRawShapeSchema,
3334
promptArgumentsFromStandardSchema,
3435
ProtocolError,
3536
ProtocolErrorCode,
@@ -873,6 +874,31 @@ export class McpServer {
873874
_meta?: Record<string, unknown>;
874875
},
875876
cb: ToolCallback<InputArgs>
877+
): RegisteredTool;
878+
/** @deprecated Wrap with `z.object({...})` instead. Raw-shape form: `inputSchema`/`outputSchema` may be a plain `{ field: z.string() }` record; it is auto-wrapped with `z.object()`. */
879+
registerTool<InputArgs extends ZodRawShape, OutputArgs extends ZodRawShape | StandardSchemaWithJSON | undefined = undefined>(
880+
name: string,
881+
config: {
882+
title?: string;
883+
description?: string;
884+
inputSchema?: InputArgs;
885+
outputSchema?: OutputArgs;
886+
annotations?: ToolAnnotations;
887+
_meta?: Record<string, unknown>;
888+
},
889+
cb: LegacyToolCallback<InputArgs>
890+
): RegisteredTool;
891+
registerTool(
892+
name: string,
893+
config: {
894+
title?: string;
895+
description?: string;
896+
inputSchema?: StandardSchemaWithJSON | ZodRawShape;
897+
outputSchema?: StandardSchemaWithJSON | ZodRawShape;
898+
annotations?: ToolAnnotations;
899+
_meta?: Record<string, unknown>;
900+
},
901+
cb: ToolCallback<StandardSchemaWithJSON | undefined> | LegacyToolCallback<ZodRawShape>
876902
): RegisteredTool {
877903
if (this._registeredTools[name]) {
878904
throw new Error(`Tool ${name} is already registered`);
@@ -884,8 +910,8 @@ export class McpServer {
884910
name,
885911
title,
886912
description,
887-
inputSchema,
888-
outputSchema,
913+
normalizeRawShapeSchema(inputSchema),
914+
normalizeRawShapeSchema(outputSchema),
889915
annotations,
890916
{ taskSupport: 'forbidden' },
891917
_meta,
@@ -928,6 +954,27 @@ export class McpServer {
928954
_meta?: Record<string, unknown>;
929955
},
930956
cb: PromptCallback<Args>
957+
): RegisteredPrompt;
958+
/** @deprecated Wrap with `z.object({...})` instead. Raw-shape form: `argsSchema` may be a plain `{ field: z.string() }` record; it is auto-wrapped with `z.object()`. */
959+
registerPrompt<Args extends ZodRawShape>(
960+
name: string,
961+
config: {
962+
title?: string;
963+
description?: string;
964+
argsSchema?: Args;
965+
_meta?: Record<string, unknown>;
966+
},
967+
cb: LegacyPromptCallback<Args>
968+
): RegisteredPrompt;
969+
registerPrompt(
970+
name: string,
971+
config: {
972+
title?: string;
973+
description?: string;
974+
argsSchema?: StandardSchemaWithJSON | ZodRawShape;
975+
_meta?: Record<string, unknown>;
976+
},
977+
cb: PromptCallback<StandardSchemaWithJSON> | LegacyPromptCallback<ZodRawShape>
931978
): RegisteredPrompt {
932979
if (this._registeredPrompts[name]) {
933980
throw new Error(`Prompt ${name} is already registered`);
@@ -939,7 +986,7 @@ export class McpServer {
939986
name,
940987
title,
941988
description,
942-
argsSchema,
989+
normalizeRawShapeSchema(argsSchema),
943990
cb as PromptCallback<StandardSchemaWithJSON | undefined>,
944991
_meta
945992
);
@@ -1062,6 +1109,25 @@ export class ResourceTemplate {
10621109
}
10631110
}
10641111

1112+
/**
1113+
* A plain record of field schemas, e.g. `{ name: z.string() }`. Accepted by
1114+
* `registerTool`/`registerPrompt` as a shorthand; auto-wrapped with `z.object()`.
1115+
*/
1116+
export type ZodRawShape = Record<string, StandardSchemaWithJSON>;
1117+
1118+
/** Infers the parsed-output type of a {@linkcode ZodRawShape}. */
1119+
export type InferRawShape<S extends ZodRawShape> = { [K in keyof S]: StandardSchemaWithJSON.InferOutput<S[K]> };
1120+
1121+
/** {@linkcode ToolCallback} variant used when `inputSchema` is a {@linkcode ZodRawShape}. */
1122+
export type LegacyToolCallback<Args extends ZodRawShape | undefined> = Args extends ZodRawShape
1123+
? (args: InferRawShape<Args>, ctx: ServerContext) => CallToolResult | Promise<CallToolResult>
1124+
: (ctx: ServerContext) => CallToolResult | Promise<CallToolResult>;
1125+
1126+
/** {@linkcode PromptCallback} variant used when `argsSchema` is a {@linkcode ZodRawShape}. */
1127+
export type LegacyPromptCallback<Args extends ZodRawShape | undefined> = Args extends ZodRawShape
1128+
? (args: InferRawShape<Args>, ctx: ServerContext) => GetPromptResult | Promise<GetPromptResult>
1129+
: (ctx: ServerContext) => GetPromptResult | Promise<GetPromptResult>;
1130+
10651131
export type BaseToolCallback<
10661132
SendResultT extends Result,
10671133
Ctx extends ServerContext,
Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
import { isStandardSchema } from '@modelcontextprotocol/core';
2+
import { describe, expect, it, vi } from 'vitest';
3+
import * as z from 'zod/v4';
4+
import { McpServer } from '../../src/index.js';
5+
6+
describe('registerTool/registerPrompt accept raw Zod shape (auto-wrapped)', () => {
7+
it('registerTool accepts a raw shape for inputSchema, auto-wraps, and does not warn', () => {
8+
const warn = vi.spyOn(console, 'warn').mockImplementation(() => {});
9+
const server = new McpServer({ name: 't', version: '1.0.0' });
10+
11+
server.registerTool('a', { inputSchema: { x: z.number() } }, async ({ x }) => ({
12+
content: [{ type: 'text' as const, text: String(x) }]
13+
}));
14+
server.registerTool('b', { inputSchema: { y: z.number() } }, async ({ y }) => ({
15+
content: [{ type: 'text' as const, text: String(y) }]
16+
}));
17+
18+
const tools = (server as unknown as { _registeredTools: Record<string, { inputSchema?: unknown }> })._registeredTools;
19+
expect(Object.keys(tools)).toEqual(['a', 'b']);
20+
// raw shape was wrapped into a Standard Schema (z.object)
21+
expect(isStandardSchema(tools['a']?.inputSchema)).toBe(true);
22+
23+
expect(warn).not.toHaveBeenCalled();
24+
warn.mockRestore();
25+
});
26+
27+
it('registerTool with z.object() inputSchema also works without warning', () => {
28+
const warn = vi.spyOn(console, 'warn').mockImplementation(() => {});
29+
const server = new McpServer({ name: 't', version: '1.0.0' });
30+
31+
server.registerTool('c', { inputSchema: z.object({ x: z.number() }) }, async ({ x }) => ({
32+
content: [{ type: 'text' as const, text: String(x) }]
33+
}));
34+
35+
expect(warn).not.toHaveBeenCalled();
36+
warn.mockRestore();
37+
});
38+
39+
it('registerPrompt accepts a raw shape for argsSchema and does not warn', () => {
40+
const warn = vi.spyOn(console, 'warn').mockImplementation(() => {});
41+
const server = new McpServer({ name: 't', version: '1.0.0' });
42+
43+
server.registerPrompt('p', { argsSchema: { topic: z.string() } }, async ({ topic }) => ({
44+
messages: [{ role: 'user' as const, content: { type: 'text' as const, text: topic } }]
45+
}));
46+
47+
const prompts = (server as unknown as { _registeredPrompts: Record<string, { argsSchema?: unknown }> })._registeredPrompts;
48+
expect(Object.keys(prompts)).toContain('p');
49+
expect(isStandardSchema(prompts['p']?.argsSchema)).toBe(true);
50+
51+
expect(warn).not.toHaveBeenCalled();
52+
warn.mockRestore();
53+
});
54+
});

0 commit comments

Comments
 (0)