Skip to content

Commit 526b97d

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

File tree

2 files changed

+182
-1
lines changed

2 files changed

+182
-1
lines changed

src/server/mcp.ts

Lines changed: 63 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -948,6 +948,13 @@ 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>(name: string, paramsSchema: Schema, cb: ToolCallback<Schema>): RegisteredTool;
957+
951958
/**
952959
* Registers a tool `name` (with a description) taking either parameter schema or annotations.
953960
* This unified overload handles both `tool(name, description, paramsSchema, cb)` and
@@ -964,6 +971,13 @@ export class McpServer {
964971
cb: ToolCallback<Args>
965972
): RegisteredTool;
966973

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

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

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

test/server/mcp.test.ts

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

0 commit comments

Comments
 (0)