Skip to content

Commit bb99b72

Browse files
committed
fix: support z.object() schemas in tool() method
1 parent 384311b commit bb99b72

File tree

2 files changed

+167
-1
lines changed

2 files changed

+167
-1
lines changed

src/server/mcp.ts

Lines changed: 66 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -948,6 +948,17 @@ export class McpServer {
948948
cb: ToolCallback<Args>
949949
): RegisteredTool;
950950

951+
/**
952+
* Registers a tool with a ZodObject schema (e.g., z.object({ ... })).
953+
* The schema's shape will be extracted and used for validation.
954+
* @deprecated Use `registerTool` instead.
955+
*/
956+
tool<Schema extends AnyObjectSchema>(
957+
name: string,
958+
paramsSchema: Schema,
959+
cb: ToolCallback<Schema>
960+
): RegisteredTool;
961+
951962
/**
952963
* Registers a tool `name` (with a description) taking either parameter schema or annotations.
953964
* This unified overload handles both `tool(name, description, paramsSchema, cb)` and
@@ -964,6 +975,18 @@ export class McpServer {
964975
cb: ToolCallback<Args>
965976
): RegisteredTool;
966977

978+
/**
979+
* Registers a tool with a description and ZodObject schema (e.g., z.object({ ... })).
980+
* The schema's shape will be extracted and used for validation.
981+
* @deprecated Use `registerTool` instead.
982+
*/
983+
tool<Schema extends AnyObjectSchema>(
984+
name: string,
985+
description: string,
986+
paramsSchema: Schema,
987+
cb: ToolCallback<Schema>
988+
): RegisteredTool;
989+
967990
/**
968991
* Registers a tool with both parameter schema and annotations.
969992
* @deprecated Use `registerTool` instead.
@@ -975,6 +998,18 @@ export class McpServer {
975998
cb: ToolCallback<Args>
976999
): RegisteredTool;
9771000

1001+
/**
1002+
* Registers a tool with a ZodObject schema and annotations.
1003+
* The schema's shape will be extracted and used for validation.
1004+
* @deprecated Use `registerTool` instead.
1005+
*/
1006+
tool<Schema extends AnyObjectSchema>(
1007+
name: string,
1008+
paramsSchema: Schema,
1009+
annotations: ToolAnnotations,
1010+
cb: ToolCallback<Schema>
1011+
): RegisteredTool;
1012+
9781013
/**
9791014
* Registers a tool with description, parameter schema, and annotations.
9801015
* @deprecated Use `registerTool` instead.
@@ -987,6 +1022,19 @@ export class McpServer {
9871022
cb: ToolCallback<Args>
9881023
): RegisteredTool;
9891024

1025+
/**
1026+
* Registers a tool with description, ZodObject schema, and annotations.
1027+
* The schema's shape will be extracted and used for validation.
1028+
* @deprecated Use `registerTool` instead.
1029+
*/
1030+
tool<Schema extends AnyObjectSchema>(
1031+
name: string,
1032+
description: string,
1033+
paramsSchema: Schema,
1034+
annotations: ToolAnnotations,
1035+
cb: ToolCallback<Schema>
1036+
): RegisteredTool;
1037+
9901038
/**
9911039
* tool() implementation. Parses arguments passed to overrides defined above.
9921040
*/
@@ -1023,8 +1071,25 @@ export class McpServer {
10231071
// Or: tool(name, description, paramsSchema, annotations, cb)
10241072
annotations = rest.shift() as ToolAnnotations;
10251073
}
1074+
} else if (typeof firstArg === 'object' && firstArg !== null && isZodSchemaInstance(firstArg)) {
1075+
// It's a Zod schema instance (like z.object()), extract its shape if it's an object schema
1076+
const shape = getObjectShape(firstArg as AnyObjectSchema);
1077+
if (shape) {
1078+
// We found an object schema, use its shape
1079+
inputSchema = shape;
1080+
rest.shift();
1081+
1082+
// Check if the next arg is potentially annotations
1083+
if (rest.length > 1 && typeof rest[0] === 'object' && rest[0] !== null && !isZodRawShapeCompat(rest[0]) && !isZodSchemaInstance(rest[0])) {
1084+
annotations = rest.shift() as ToolAnnotations;
1085+
}
1086+
} else {
1087+
// It's a schema but not an object schema, treat as annotations
1088+
// (This maintains backward compatibility for edge cases)
1089+
annotations = rest.shift() as ToolAnnotations;
1090+
}
10261091
} else if (typeof firstArg === 'object' && firstArg !== null) {
1027-
// Not a ZodRawShapeCompat, so must be annotations in this position
1092+
// Not a ZodRawShapeCompat or Zod schema, so must be annotations in this position
10281093
// Case: tool(name, annotations, cb)
10291094
// Or: tool(name, description, annotations, cb)
10301095
annotations = rest.shift() as ToolAnnotations;

test/server/mcp.test.ts

Lines changed: 101 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -993,6 +993,107 @@ describe.each(zodTestMatrix)('$zodVersionLabel', (entry: ZodMatrixEntry) => {
993993
expect(result.tools[1].annotations).toEqual(result.tools[0].annotations);
994994
});
995995

996+
test('should accept z.object() schemas and extract arguments correctly', async () => {
997+
const mcpServer = new McpServer({
998+
name: 'test server',
999+
version: '1.0'
1000+
});
1001+
const client = new Client({
1002+
name: 'test client',
1003+
version: '1.0'
1004+
});
1005+
1006+
// tool() with z.object() schema
1007+
mcpServer.tool(
1008+
'test-zobject',
1009+
'Test with ZodObject',
1010+
z.object({
1011+
message: z.string()
1012+
}),
1013+
async ({ message }) => ({
1014+
content: [{ type: 'text', text: `Echo: ${message}` }]
1015+
})
1016+
);
1017+
1018+
// tool() with z.object() and annotations
1019+
mcpServer.tool(
1020+
'test-zobject-annotations',
1021+
'Test with ZodObject and annotations',
1022+
z.object({
1023+
name: z.string(),
1024+
value: z.number()
1025+
}),
1026+
{ title: 'ZodObject Tool', readOnlyHint: true },
1027+
async ({ name, value }) => ({
1028+
content: [{ type: 'text', text: `${name}: ${value}` }]
1029+
})
1030+
);
1031+
1032+
const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair();
1033+
1034+
await Promise.all([client.connect(clientTransport), mcpServer.connect(serverTransport)]);
1035+
1036+
const listResult = await client.request({ method: 'tools/list' }, ListToolsResultSchema);
1037+
1038+
expect(listResult.tools).toHaveLength(2);
1039+
expect(listResult.tools[0].name).toBe('test-zobject');
1040+
expect(listResult.tools[0].inputSchema).toMatchObject({
1041+
type: 'object',
1042+
properties: {
1043+
message: { type: 'string' }
1044+
}
1045+
});
1046+
1047+
expect(listResult.tools[1].name).toBe('test-zobject-annotations');
1048+
expect(listResult.tools[1].inputSchema).toMatchObject({
1049+
type: 'object',
1050+
properties: {
1051+
name: { type: 'string' },
1052+
value: { type: 'number' }
1053+
}
1054+
});
1055+
expect(listResult.tools[1].annotations).toEqual({
1056+
title: 'ZodObject Tool',
1057+
readOnlyHint: true
1058+
});
1059+
1060+
const result1 = await client.request(
1061+
{
1062+
method: 'tools/call',
1063+
params: {
1064+
name: 'test-zobject',
1065+
arguments: { message: 'Hello World' }
1066+
}
1067+
},
1068+
CallToolResultSchema
1069+
);
1070+
1071+
expect(result1.content).toEqual([
1072+
{
1073+
type: 'text',
1074+
text: 'Echo: Hello World'
1075+
}
1076+
]);
1077+
1078+
const result2 = await client.request(
1079+
{
1080+
method: 'tools/call',
1081+
params: {
1082+
name: 'test-zobject-annotations',
1083+
arguments: { name: 'test', value: 42 }
1084+
}
1085+
},
1086+
CallToolResultSchema
1087+
);
1088+
1089+
expect(result2.content).toEqual([
1090+
{
1091+
type: 'text',
1092+
text: 'test: 42'
1093+
}
1094+
]);
1095+
});
1096+
9961097
/***
9971098
* Test: Tool Argument Validation
9981099
*/

0 commit comments

Comments
 (0)