Skip to content

Support ZodObject schemas in tool method #1310

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
DaleSeo wants to merge 1 commit into
base: main
Choose a base branch
from
+182 −1
Draft

Support ZodObject schemas in tool method #1310

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
64 changes: 63 additions & 1 deletion src/server/mcp.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -948,6 +948,13 @@ export class McpServer {
cb: ToolCallback<Args>
): RegisteredTool;

/**
* Registers a tool with a ZodObject schema (e.g., z.object({ ... })).
* The schema's shape will be extracted and used for validation.
* @deprecated Use `registerTool` instead.
*/
tool<Schema extends AnyObjectSchema>(name: string, paramsSchema: Schema, cb: ToolCallback<Schema>): RegisteredTool;

/**
* Registers a tool `name` (with a description) taking either parameter schema or annotations.
* This unified overload handles both `tool(name, description, paramsSchema, cb)` and
Expand All @@ -964,6 +971,13 @@ export class McpServer {
cb: ToolCallback<Args>
): RegisteredTool;

/**
* Registers a tool with a description and ZodObject schema (e.g., z.object({ ... })).
* The schema's shape will be extracted and used for validation.
* @deprecated Use `registerTool` instead.
*/
tool<Schema extends AnyObjectSchema>(name: string, description: string, paramsSchema: Schema, cb: ToolCallback<Schema>): RegisteredTool;

/**
* Registers a tool with both parameter schema and annotations.
* @deprecated Use `registerTool` instead.
Expand All @@ -975,6 +989,18 @@ export class McpServer {
cb: ToolCallback<Args>
): RegisteredTool;

/**
* Registers a tool with a ZodObject schema and annotations.
* The schema's shape will be extracted and used for validation.
* @deprecated Use `registerTool` instead.
*/
tool<Schema extends AnyObjectSchema>(
name: string,
paramsSchema: Schema,
annotations: ToolAnnotations,
cb: ToolCallback<Schema>
): RegisteredTool;

/**
* Registers a tool with description, parameter schema, and annotations.
* @deprecated Use `registerTool` instead.
Expand All @@ -987,6 +1013,19 @@ export class McpServer {
cb: ToolCallback<Args>
): RegisteredTool;

/**
* Registers a tool with description, ZodObject schema, and annotations.
* The schema's shape will be extracted and used for validation.
* @deprecated Use `registerTool` instead.
*/
tool<Schema extends AnyObjectSchema>(
name: string,
description: string,
paramsSchema: Schema,
annotations: ToolAnnotations,
cb: ToolCallback<Schema>
): RegisteredTool;

/**
* tool() implementation. Parses arguments passed to overrides defined above.
*/
Expand Down Expand Up @@ -1023,8 +1062,31 @@ export class McpServer {
// Or: tool(name, description, paramsSchema, annotations, cb)
annotations = rest.shift() as ToolAnnotations;
}
} else if (typeof firstArg === 'object' && firstArg !== null && isZodSchemaInstance(firstArg)) {
// It's a Zod schema instance (like z.object()), extract its shape if it's an object schema
const shape = getObjectShape(firstArg as AnyObjectSchema);
if (shape) {
// We found an object schema, use its shape
inputSchema = shape;
rest.shift();

// Check if the next arg is potentially annotations
if (
rest.length > 1 &&
typeof rest[0] === 'object' &&
rest[0] !== null &&
!isZodRawShapeCompat(rest[0]) &&
!isZodSchemaInstance(rest[0])
) {
annotations = rest.shift() as ToolAnnotations;
}
} else {
// It's a schema but not an object schema, treat as annotations
// (This maintains backward compatibility for edge cases)
annotations = rest.shift() as ToolAnnotations;
}
} else if (typeof firstArg === 'object' && firstArg !== null) {
// Not a ZodRawShapeCompat, so must be annotations in this position
// Not a ZodRawShapeCompat or Zod schema, so must be annotations in this position
// Case: tool(name, annotations, cb)
// Or: tool(name, description, annotations, cb)
annotations = rest.shift() as ToolAnnotations;
Expand Down
119 changes: 119 additions & 0 deletions test/server/mcp.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -993,6 +993,125 @@ describe.each(zodTestMatrix)('$zodVersionLabel', (entry: ZodMatrixEntry) => {
expect(result.tools[1].annotations).toEqual(result.tools[0].annotations);
});

test('should accept z.object() schema and pass arguments correctly', async () => {
const mcpServer = new McpServer({
name: 'test server',
version: '1.0'
});
const client = new Client({
name: 'test client',
version: '1.0'
});

mcpServer.tool(
'test-zobject',
'Test with ZodObject',
z.object({
message: z.string()
}),
async ({ message }) => ({
content: [{ type: 'text', text: `Echo: ${message}` }]
})
);

const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair();

await Promise.all([client.connect(clientTransport), mcpServer.connect(serverTransport)]);

const listResult = await client.request({ method: 'tools/list' }, ListToolsResultSchema);

expect(listResult.tools).toHaveLength(1);
expect(listResult.tools[0].name).toBe('test-zobject');
expect(listResult.tools[0].inputSchema).toMatchObject({
type: 'object',
properties: {
message: { type: 'string' }
}
});

// Verify tool call receives arguments correctly (not the extra object)
const result = await client.request(
{
method: 'tools/call',
params: {
name: 'test-zobject',
arguments: { message: 'Hello World' }
}
},
CallToolResultSchema
);

expect(result.content).toEqual([
{
type: 'text',
text: 'Echo: Hello World'
}
]);
});

test('should accept z.object() schema with annotations', async () => {
const mcpServer = new McpServer({
name: 'test server',
version: '1.0'
});
const client = new Client({
name: 'test client',
version: '1.0'
});

mcpServer.tool(
'test-zobject-annotations',
'Test with ZodObject and annotations',
z.object({
name: z.string(),
value: z.number()
}),
{ title: 'ZodObject Tool', readOnlyHint: true },
async ({ name, value }) => ({
content: [{ type: 'text', text: `${name}: ${value}` }]
})
);

const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair();

await Promise.all([client.connect(clientTransport), mcpServer.connect(serverTransport)]);

const listResult = await client.request({ method: 'tools/list' }, ListToolsResultSchema);

expect(listResult.tools).toHaveLength(1);
expect(listResult.tools[0].name).toBe('test-zobject-annotations');
expect(listResult.tools[0].inputSchema).toMatchObject({
type: 'object',
properties: {
name: { type: 'string' },
value: { type: 'number' }
}
});
expect(listResult.tools[0].annotations).toEqual({
title: 'ZodObject Tool',
readOnlyHint: true
});

// Verify tool call receives arguments correctly
const result = await client.request(
{
method: 'tools/call',
params: {
name: 'test-zobject-annotations',
arguments: { name: 'test', value: 42 }
}
},
CallToolResultSchema
);

expect(result.content).toEqual([
{
type: 'text',
text: 'test: 42'
}
]);
});

/***
* Test: Tool Argument Validation
*/
Expand Down
Loading