Skip to content

Commit 9bf5ad5

Browse files
DaleSeoDaleSeo
committed
fix: support z.object() schemas in tool() method
1 parent f82c997 commit 9bf5ad5

File tree

2 files changed

+182
-1
lines changed

2 files changed

+182
-1
lines changed

‎packages/server/src/server/mcp.ts‎

Lines changed: 63 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -950,6 +950,13 @@ export class McpServer {
950950
cb: ToolCallback<Args>
951951
): RegisteredTool;
952952

953+
/**
954+
* Registers a tool with a ZodObject schema (e.g., z.object({ ... })).
955+
* The schema's shape will be extracted and used for validation.
956+
* @deprecated Use `registerTool` instead.
957+
*/
958+
tool<Schema extends AnyObjectSchema>(name: string, paramsSchema: Schema, cb: ToolCallback<Schema>): RegisteredTool;
959+
953960
/**
954961
* Registers a tool `name` (with a description) taking either parameter schema or annotations.
955962
* This unified overload handles both `tool(name, description, paramsSchema, cb)` and
@@ -966,6 +973,13 @@ export class McpServer {
966973
cb: ToolCallback<Args>
967974
): RegisteredTool;
968975

976+
/**
977+
* Registers a tool with a description and ZodObject schema (e.g., z.object({ ... })).
978+
* The schema's shape will be extracted and used for validation.
979+
* @deprecated Use `registerTool` instead.
980+
*/
981+
tool<Schema extends AnyObjectSchema>(name: string, description: string, paramsSchema: Schema, cb: ToolCallback<Schema>): RegisteredTool;
982+
969983
/**
970984
* Registers a tool with both parameter schema and annotations.
971985
* @deprecated Use `registerTool` instead.
@@ -977,6 +991,18 @@ export class McpServer {
977991
cb: ToolCallback<Args>
978992
): RegisteredTool;
979993

994+
/**
995+
* Registers a tool with a ZodObject schema and annotations.
996+
* The schema's shape will be extracted and used for validation.
997+
* @deprecated Use `registerTool` instead.
998+
*/
999+
tool<Schema extends AnyObjectSchema>(
1000+
name: string,
1001+
paramsSchema: Schema,
1002+
annotations: ToolAnnotations,
1003+
cb: ToolCallback<Schema>
1004+
): RegisteredTool;
1005+
9801006
/**
9811007
* Registers a tool with description, parameter schema, and annotations.
9821008
* @deprecated Use `registerTool` instead.
@@ -989,6 +1015,19 @@ export class McpServer {
9891015
cb: ToolCallback<Args>
9901016
): RegisteredTool;
9911017

1018+
/**
1019+
* Registers a tool with description, ZodObject schema, and annotations.
1020+
* The schema's shape will be extracted and used for validation.
1021+
* @deprecated Use `registerTool` instead.
1022+
*/
1023+
tool<Schema extends AnyObjectSchema>(
1024+
name: string,
1025+
description: string,
1026+
paramsSchema: Schema,
1027+
annotations: ToolAnnotations,
1028+
cb: ToolCallback<Schema>
1029+
): RegisteredTool;
1030+
9921031
/**
9931032
* tool() implementation. Parses arguments passed to overrides defined above.
9941033
*/
@@ -1025,8 +1064,31 @@ export class McpServer {
10251064
// Or: tool(name, description, paramsSchema, annotations, cb)
10261065
annotations = rest.shift() as ToolAnnotations;
10271066
}
1067+
} else if (typeof firstArg === 'object' && firstArg !== null && isZodSchemaInstance(firstArg)) {
1068+
// It's a Zod schema instance (like z.object()), extract its shape if it's an object schema
1069+
const shape = getObjectShape(firstArg as AnyObjectSchema);
1070+
if (shape) {
1071+
// We found an object schema, use its shape
1072+
inputSchema = shape;
1073+
rest.shift();
1074+
1075+
// Check if the next arg is potentially annotations
1076+
if (
1077+
rest.length > 1 &&
1078+
typeof rest[0] === 'object' &&
1079+
rest[0] !== null &&
1080+
!isZodRawShapeCompat(rest[0]) &&
1081+
!isZodSchemaInstance(rest[0])
1082+
) {
1083+
annotations = rest.shift() as ToolAnnotations;
1084+
}
1085+
} else {
1086+
// It's a schema but not an object schema, treat as annotations
1087+
// (This maintains backward compatibility for edge cases)
1088+
annotations = rest.shift() as ToolAnnotations;
1089+
}
10281090
} else if (typeof firstArg === 'object' && firstArg !== null) {
1029-
// Not a ZodRawShapeCompat, so must be annotations in this position
1091+
// Not a ZodRawShapeCompat or Zod schema, so must be annotations in this position
10301092
// Case: tool(name, annotations, cb)
10311093
// Or: tool(name, description, annotations, cb)
10321094
annotations = rest.shift() as ToolAnnotations;

‎test/integration/test/server/mcp.test.ts‎

Lines changed: 119 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -991,6 +991,125 @@ describe.each(zodTestMatrix)('$zodVersionLabel', (entry: ZodMatrixEntry) => {
991991
expect(result.tools[1]!.annotations).toEqual(result.tools[0]!.annotations);
992992
});
993993

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

0 commit comments

Comments
 (0)