Skip to content

Commit e15a8ef

Browse files
feat(compat): registerTool/registerPrompt accept raw Zod shape, auto-wrap with z.object() (#1901)
Co-authored-by: Konstantin Konstantinov <KKonstantinov@users.noreply.github.com>
1 parent 9fc9070 commit e15a8ef

7 files changed

Lines changed: 387 additions & 12 deletions

File tree

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
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()`. The raw-shape overloads are `@deprecated`; prefer wrapping with `z.object()`.
7+
8+
Also widens the `completable()` constraint from `StandardSchemaWithJSON` to `StandardSchemaV1` so v1's `completable(z.string(), fn)` continues to work.

packages/core/src/index.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ export * from './types/index.js';
1515
export * from './util/inMemory.js';
1616
export * from './util/schema.js';
1717
export * from './util/standardSchema.js';
18+
export * from './util/zodCompat.js';
1819

1920
// experimental exports
2021
export * from './experimental/index.js';
Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,80 @@
1+
/**
2+
* Zod-specific helpers for the v1-compat raw-shape shorthand on
3+
* `registerTool`/`registerPrompt`. Kept separate from `standardSchema.ts` so
4+
* that file stays library-agnostic per the Standard Schema spec.
5+
*/
6+
7+
import * as z from 'zod/v4';
8+
9+
import type { StandardSchemaWithJSON } from './standardSchema.js';
10+
import { isStandardSchema } from './standardSchema.js';
11+
12+
function isZodV4Schema(v: unknown): v is z.ZodType {
13+
// `_zod` is the v4 internal namespace property. Zod v3 schemas have `_def`
14+
// and (since 3.24) `~standard.vendor === 'zod'`, but never `_zod`. We require
15+
// v4 because the wrap path below uses v4's `z.object()`, which cannot consume
16+
// v3 field schemas.
17+
return typeof v === 'object' && v !== null && '_zod' in v;
18+
}
19+
20+
function looksLikeZodV3(v: unknown): boolean {
21+
// v3 schemas have `_def.typeName` (e.g. 'ZodString') and no `_zod`.
22+
return (
23+
typeof v === 'object' &&
24+
v !== null &&
25+
!('_zod' in v) &&
26+
'_def' in v &&
27+
typeof (v as { _def?: { typeName?: unknown } })._def?.typeName === 'string'
28+
);
29+
}
30+
31+
/**
32+
* Detects a "raw shape" — a plain object whose values are Zod field schemas,
33+
* e.g. `{ name: z.string() }`. Powers the auto-wrap in
34+
* {@linkcode normalizeRawShapeSchema}, which wraps with `z.object()`, so only
35+
* Zod values are supported.
36+
*
37+
* @internal
38+
*/
39+
export function isZodRawShape(obj: unknown): obj is Record<string, z.ZodType> {
40+
if (typeof obj !== 'object' || obj === null) return false;
41+
if (isStandardSchema(obj)) return false;
42+
// Require a plain object literal: rejects arrays, Date, Map, RegExp, class instances, etc.
43+
// Object.create(null) is also accepted.
44+
const proto = Object.getPrototypeOf(obj);
45+
if (proto !== Object.prototype && proto !== null) return false;
46+
// [].every() is true, so an empty plain object is a valid raw shape (matches v1).
47+
return Object.values(obj).every(v => isZodV4Schema(v));
48+
}
49+
50+
/**
51+
* Accepts either a {@linkcode StandardSchemaWithJSON} or a raw Zod shape
52+
* `{ field: z.string() }` and returns a {@linkcode StandardSchemaWithJSON}.
53+
* Raw shapes are wrapped with `z.object()` so the rest of the pipeline sees a
54+
* uniform schema type; already-wrapped schemas pass through unchanged.
55+
*
56+
* @internal
57+
*/
58+
export function normalizeRawShapeSchema(
59+
schema: StandardSchemaWithJSON | Record<string, z.ZodType> | undefined
60+
): StandardSchemaWithJSON | undefined {
61+
if (schema === undefined) return undefined;
62+
if (isZodRawShape(schema)) {
63+
return z.object(schema) as StandardSchemaWithJSON;
64+
}
65+
if (typeof schema === 'object' && schema !== null && !isStandardSchema(schema) && Object.values(schema).some(v => looksLikeZodV3(v))) {
66+
throw new TypeError(
67+
'Raw-shape inputSchema/outputSchema/argsSchema fields must be Zod v4 schemas. Got a Zod v3 field schema. Import from `zod/v4` (or upgrade your zod import), or wrap with `z.object({...})` yourself.'
68+
);
69+
}
70+
if (!isStandardSchema(schema)) {
71+
throw new TypeError(
72+
'inputSchema/outputSchema/argsSchema must be a Standard Schema (e.g. z.object({...})) or a raw Zod shape ({ field: z.string() }).'
73+
);
74+
}
75+
// Any StandardSchema passes through; standardSchemaToJsonSchema owns the per-vendor
76+
// handling for schemas without `~standard.jsonSchema` (zod 4.0-4.1 fallback, zod 3
77+
// and non-zod errors). Gating on `~standard.jsonSchema` here would unreachably
78+
// front-run that fallback.
79+
return schema;
80+
}
Lines changed: 89 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,89 @@
1+
import { vi } from 'vitest';
2+
import * as z from 'zod/v4';
3+
4+
import { standardSchemaToJsonSchema } from '../../src/util/standardSchema.js';
5+
import { isZodRawShape, normalizeRawShapeSchema } from '../../src/util/zodCompat.js';
6+
7+
describe('isZodRawShape', () => {
8+
test('treats empty object as a raw shape (matches v1)', () => {
9+
expect(isZodRawShape({})).toBe(true);
10+
});
11+
test('detects raw shape with zod fields', () => {
12+
expect(isZodRawShape({ a: z.string() })).toBe(true);
13+
});
14+
test('rejects a Standard Schema instance', () => {
15+
expect(isZodRawShape(z.object({ a: z.string() }))).toBe(false);
16+
});
17+
test('rejects a shape with non-Zod Standard Schema fields', () => {
18+
const nonZod = { '~standard': { version: 1, vendor: 'arktype', validate: () => ({ value: 'x' }) } };
19+
expect(isZodRawShape({ a: nonZod })).toBe(false);
20+
});
21+
test('rejects a shape with Zod v3 fields (only v4 is wrappable)', () => {
22+
expect(isZodRawShape({ a: mockZodV3String() })).toBe(false);
23+
});
24+
test('rejects non-plain objects with no own-enumerable properties', () => {
25+
expect(isZodRawShape([])).toBe(false);
26+
expect(isZodRawShape([z.string()])).toBe(false);
27+
expect(isZodRawShape(new Date())).toBe(false);
28+
expect(isZodRawShape(new Map())).toBe(false);
29+
expect(isZodRawShape(/regex/)).toBe(false);
30+
});
31+
test('accepts a null-prototype plain object', () => {
32+
const o = Object.create(null);
33+
o.a = z.string();
34+
expect(isZodRawShape(o)).toBe(true);
35+
});
36+
});
37+
38+
// Minimal structural mock of a Zod v3 schema: has `_def.typeName` and
39+
// `~standard.vendor === 'zod'` (zod >=3.24), but no `_zod`.
40+
function mockZodV3String(): unknown {
41+
return {
42+
_def: { typeName: 'ZodString', checks: [], coerce: false },
43+
'~standard': { version: 1, vendor: 'zod', validate: (v: unknown) => ({ value: v }) },
44+
parse: (v: unknown) => v
45+
};
46+
}
47+
48+
describe('normalizeRawShapeSchema', () => {
49+
test('wraps empty raw shape into z.object({})', () => {
50+
const wrapped = normalizeRawShapeSchema({});
51+
expect(wrapped).toBeDefined();
52+
expect(standardSchemaToJsonSchema(wrapped!, 'input').type).toBe('object');
53+
});
54+
test('passes through an already-wrapped Standard Schema unchanged', () => {
55+
const schema = z.object({ a: z.string() });
56+
expect(normalizeRawShapeSchema(schema)).toBe(schema);
57+
});
58+
test('returns undefined for undefined input', () => {
59+
expect(normalizeRawShapeSchema(undefined)).toBeUndefined();
60+
});
61+
test('throws TypeError for an invalid object that is neither raw shape nor Standard Schema', () => {
62+
expect(() => normalizeRawShapeSchema({ a: 'not a zod schema' } as never)).toThrow(TypeError);
63+
});
64+
test('passes through a Standard Schema without `~standard.jsonSchema` (per-vendor handling deferred to standardSchemaToJsonSchema)', () => {
65+
const noJson = { '~standard': { version: 1, vendor: 'x', validate: () => ({ value: {} }) } };
66+
expect(normalizeRawShapeSchema(noJson as never)).toBe(noJson);
67+
});
68+
test('passes through a zod 4.0-4.1 schema so standardSchemaToJsonSchema can apply its z.toJSONSchema fallback', () => {
69+
const warn = vi.spyOn(console, 'warn').mockImplementation(() => {});
70+
const real = z.object({ a: z.string() });
71+
// Simulate zod 4.0-4.1: shadow `~standard` with `jsonSchema` removed, keep `_zod` intact.
72+
const { jsonSchema: _drop, ...stdNoJson } = real['~standard'] as unknown as Record<string, unknown>;
73+
void _drop;
74+
Object.defineProperty(real, '~standard', { value: { ...stdNoJson, vendor: 'zod' }, configurable: true });
75+
76+
const normalized = normalizeRawShapeSchema(real);
77+
expect(normalized).toBe(real);
78+
const json = standardSchemaToJsonSchema(normalized!, 'input');
79+
expect(json.type).toBe('object');
80+
expect((json.properties as Record<string, unknown>)?.a).toBeDefined();
81+
warn.mockRestore();
82+
});
83+
test('throws actionable TypeError for a raw shape with Zod v3 fields', () => {
84+
expect(() => normalizeRawShapeSchema({ a: mockZodV3String() } as never)).toThrow(/Zod v4 schemas.*Got a Zod v3 field schema/);
85+
});
86+
test('throws the intended TypeError (not Object.values crash) for null input', () => {
87+
expect(() => normalizeRawShapeSchema(null as never)).toThrow(/must be a Standard Schema/);
88+
});
89+
});

packages/server/src/server/completable.ts

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,19 +1,19 @@
1-
import type { StandardSchemaWithJSON } from '@modelcontextprotocol/core';
1+
import type { StandardSchemaV1 } from '@modelcontextprotocol/core';
22

33
export const COMPLETABLE_SYMBOL: unique symbol = Symbol.for('mcp.completable');
44

5-
export type CompleteCallback<T extends StandardSchemaWithJSON = StandardSchemaWithJSON> = (
6-
value: StandardSchemaWithJSON.InferInput<T>,
5+
export type CompleteCallback<T extends StandardSchemaV1 = StandardSchemaV1> = (
6+
value: StandardSchemaV1.InferInput<T>,
77
context?: {
88
arguments?: Record<string, string>;
99
}
10-
) => StandardSchemaWithJSON.InferInput<T>[] | Promise<StandardSchemaWithJSON.InferInput<T>[]>;
10+
) => StandardSchemaV1.InferInput<T>[] | Promise<StandardSchemaV1.InferInput<T>[]>;
1111

12-
export type CompletableMeta<T extends StandardSchemaWithJSON = StandardSchemaWithJSON> = {
12+
export type CompletableMeta<T extends StandardSchemaV1 = StandardSchemaV1> = {
1313
complete: CompleteCallback<T>;
1414
};
1515

16-
export type CompletableSchema<T extends StandardSchemaWithJSON> = T & {
16+
export type CompletableSchema<T extends StandardSchemaV1> = T & {
1717
[COMPLETABLE_SYMBOL]: CompletableMeta<T>;
1818
};
1919

@@ -48,7 +48,7 @@ export type CompletableSchema<T extends StandardSchemaWithJSON> = T & {
4848
*
4949
* @see {@linkcode server/mcp.McpServer.registerPrompt | McpServer.registerPrompt} for using completable schemas in prompt argument definitions
5050
*/
51-
export function completable<T extends StandardSchemaWithJSON>(schema: T, complete: CompleteCallback<T>): CompletableSchema<T> {
51+
export function completable<T extends StandardSchemaV1>(schema: T, complete: CompleteCallback<T>): CompletableSchema<T> {
5252
Object.defineProperty(schema as object, COMPLETABLE_SYMBOL, {
5353
value: { complete } as CompletableMeta<T>,
5454
enumerable: false,
@@ -61,14 +61,14 @@ export function completable<T extends StandardSchemaWithJSON>(schema: T, complet
6161
/**
6262
* Checks if a schema is completable (has completion metadata).
6363
*/
64-
export function isCompletable(schema: unknown): schema is CompletableSchema<StandardSchemaWithJSON> {
64+
export function isCompletable(schema: unknown): schema is CompletableSchema<StandardSchemaV1> {
6565
return !!schema && typeof schema === 'object' && COMPLETABLE_SYMBOL in (schema as object);
6666
}
6767

6868
/**
6969
* Gets the completer callback from a completable schema, if it exists.
7070
*/
71-
export function getCompleter<T extends StandardSchemaWithJSON>(schema: T): CompleteCallback<T> | undefined {
71+
export function getCompleter<T extends StandardSchemaV1>(schema: T): CompleteCallback<T> | undefined {
7272
const meta = (schema as unknown as { [COMPLETABLE_SYMBOL]?: CompletableMeta<T> })[COMPLETABLE_SYMBOL];
7373
return meta?.complete as CompleteCallback<T> | undefined;
7474
}

packages/server/src/server/mcp.ts

Lines changed: 71 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,
@@ -38,6 +39,7 @@ import {
3839
validateAndWarnToolName,
3940
validateStandardSchema
4041
} from '@modelcontextprotocol/core';
42+
import type * as z from 'zod/v4';
4143

4244
import type { ToolTaskHandler } from '../experimental/tasks/interfaces.js';
4345
import { ExperimentalMcpServerTasks } from '../experimental/tasks/mcpServer.js';
@@ -873,6 +875,31 @@ export class McpServer {
873875
_meta?: Record<string, unknown>;
874876
},
875877
cb: ToolCallback<InputArgs>
878+
): RegisteredTool;
879+
/** @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()`. */
880+
registerTool<InputArgs extends ZodRawShape, OutputArgs extends ZodRawShape | StandardSchemaWithJSON | undefined = undefined>(
881+
name: string,
882+
config: {
883+
title?: string;
884+
description?: string;
885+
inputSchema?: InputArgs;
886+
outputSchema?: OutputArgs;
887+
annotations?: ToolAnnotations;
888+
_meta?: Record<string, unknown>;
889+
},
890+
cb: LegacyToolCallback<InputArgs>
891+
): RegisteredTool;
892+
registerTool(
893+
name: string,
894+
config: {
895+
title?: string;
896+
description?: string;
897+
inputSchema?: StandardSchemaWithJSON | ZodRawShape;
898+
outputSchema?: StandardSchemaWithJSON | ZodRawShape;
899+
annotations?: ToolAnnotations;
900+
_meta?: Record<string, unknown>;
901+
},
902+
cb: ToolCallback<StandardSchemaWithJSON | undefined> | LegacyToolCallback<ZodRawShape>
876903
): RegisteredTool {
877904
if (this._registeredTools[name]) {
878905
throw new Error(`Tool ${name} is already registered`);
@@ -884,8 +911,8 @@ export class McpServer {
884911
name,
885912
title,
886913
description,
887-
inputSchema,
888-
outputSchema,
914+
normalizeRawShapeSchema(inputSchema),
915+
normalizeRawShapeSchema(outputSchema),
889916
annotations,
890917
{ taskSupport: 'forbidden' },
891918
_meta,
@@ -928,6 +955,27 @@ export class McpServer {
928955
_meta?: Record<string, unknown>;
929956
},
930957
cb: PromptCallback<Args>
958+
): RegisteredPrompt;
959+
/** @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()`. */
960+
registerPrompt<Args extends ZodRawShape>(
961+
name: string,
962+
config: {
963+
title?: string;
964+
description?: string;
965+
argsSchema?: Args;
966+
_meta?: Record<string, unknown>;
967+
},
968+
cb: LegacyPromptCallback<Args>
969+
): RegisteredPrompt;
970+
registerPrompt(
971+
name: string,
972+
config: {
973+
title?: string;
974+
description?: string;
975+
argsSchema?: StandardSchemaWithJSON | ZodRawShape;
976+
_meta?: Record<string, unknown>;
977+
},
978+
cb: PromptCallback<StandardSchemaWithJSON> | LegacyPromptCallback<ZodRawShape>
931979
): RegisteredPrompt {
932980
if (this._registeredPrompts[name]) {
933981
throw new Error(`Prompt ${name} is already registered`);
@@ -939,7 +987,7 @@ export class McpServer {
939987
name,
940988
title,
941989
description,
942-
argsSchema,
990+
normalizeRawShapeSchema(argsSchema),
943991
cb as PromptCallback<StandardSchemaWithJSON | undefined>,
944992
_meta
945993
);
@@ -1062,6 +1110,26 @@ export class ResourceTemplate {
10621110
}
10631111
}
10641112

1113+
/**
1114+
* A plain record of Zod field schemas, e.g. `{ name: z.string() }`. Accepted by
1115+
* `registerTool`/`registerPrompt` as a shorthand; auto-wrapped with `z.object()`.
1116+
* Zod schemas only — `z.object()` cannot wrap other Standard Schema libraries.
1117+
*/
1118+
export type ZodRawShape = Record<string, z.ZodType>;
1119+
1120+
/** Infers the parsed-output type of a {@linkcode ZodRawShape}. */
1121+
export type InferRawShape<S extends ZodRawShape> = z.infer<z.ZodObject<S>>;
1122+
1123+
/** {@linkcode ToolCallback} variant used when `inputSchema` is a {@linkcode ZodRawShape}. */
1124+
export type LegacyToolCallback<Args extends ZodRawShape | undefined> = Args extends ZodRawShape
1125+
? (args: InferRawShape<Args>, ctx: ServerContext) => CallToolResult | Promise<CallToolResult>
1126+
: (ctx: ServerContext) => CallToolResult | Promise<CallToolResult>;
1127+
1128+
/** {@linkcode PromptCallback} variant used when `argsSchema` is a {@linkcode ZodRawShape}. */
1129+
export type LegacyPromptCallback<Args extends ZodRawShape | undefined> = Args extends ZodRawShape
1130+
? (args: InferRawShape<Args>, ctx: ServerContext) => GetPromptResult | Promise<GetPromptResult>
1131+
: (ctx: ServerContext) => GetPromptResult | Promise<GetPromptResult>;
1132+
10651133
export type BaseToolCallback<
10661134
SendResultT extends Result,
10671135
Ctx extends ServerContext,

0 commit comments

Comments
 (0)