Add client-side validation for tool output schemas
- Cache tool output schemas when listing tools - Validate structuredContent against outputSchema during callTool - Enforce that tools with outputSchema must return structuredContent - Add json-schema-to-zod dependency for schema conversion - Add comprehensive tests for validation scenarios 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
Basil Hosmer
parent
a85fee407a
commit
cf12fe0d83
4 changed files with 533 additions and 2 deletions
10
package-lock.json
generated
10
package-lock.json
generated
|
|
@ -15,6 +15,7 @@
|
|||
"eventsource": "^3.0.2",
|
||||
"express": "^5.0.1",
|
||||
"express-rate-limit": "^7.5.0",
|
||||
"json-schema-to-zod": "^2.6.1",
|
||||
"pkce-challenge": "^5.0.0",
|
||||
"raw-body": "^3.0.0",
|
||||
"zod": "^3.23.8",
|
||||
|
|
@ -4879,6 +4880,15 @@
|
|||
"integrity": "sha512-xyFwyhro/JEof6Ghe2iz2NcXoj2sloNsWr/XsERDK/oiPCfaNhl5ONfp+jQdAZRQQ0IJWNzH9zIZF7li91kh2w==",
|
||||
"dev": true
|
||||
},
|
||||
"node_modules/json-schema-to-zod": {
|
||||
"version": "2.6.1",
|
||||
"resolved": "https://registry.npmjs.org/json-schema-to-zod/-/json-schema-to-zod-2.6.1.tgz",
|
||||
"integrity": "sha512-uiHmWH21h9FjKJkRBntfVGTLpYlCZ1n98D0izIlByqQLqpmkQpNTBtfbdP04Na6+43lgsvrShFh2uWLkQDKJuQ==",
|
||||
"license": "ISC",
|
||||
"bin": {
|
||||
"json-schema-to-zod": "dist/cjs/cli.js"
|
||||
}
|
||||
},
|
||||
"node_modules/json-schema-traverse": {
|
||||
"version": "0.4.1",
|
||||
"resolved": "https://registry.npmjs.org/json-schema-traverse/-/json-schema-traverse-0.4.1.tgz",
|
||||
|
|
|
|||
|
|
@ -52,6 +52,7 @@
|
|||
"eventsource": "^3.0.2",
|
||||
"express": "^5.0.1",
|
||||
"express-rate-limit": "^7.5.0",
|
||||
"json-schema-to-zod": "^2.6.1",
|
||||
"pkce-challenge": "^5.0.0",
|
||||
"raw-body": "^3.0.0",
|
||||
"zod": "^3.23.8",
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ import {
|
|||
InitializeRequestSchema,
|
||||
ListResourcesRequestSchema,
|
||||
ListToolsRequestSchema,
|
||||
CallToolRequestSchema,
|
||||
CreateMessageRequestSchema,
|
||||
ListRootsRequestSchema,
|
||||
ErrorCode,
|
||||
|
|
@ -754,3 +755,454 @@ test("should handle request timeout", async () => {
|
|||
code: ErrorCode.RequestTimeout,
|
||||
});
|
||||
});
|
||||
|
||||
describe('outputSchema validation', () => {
|
||||
test('should validate structuredContent against outputSchema', async () => {
|
||||
const server = new Server({
|
||||
name: 'test-server',
|
||||
version: '1.0.0',
|
||||
}, {
|
||||
capabilities: {
|
||||
tools: {},
|
||||
},
|
||||
});
|
||||
|
||||
// Set up server handlers
|
||||
server.setRequestHandler(InitializeRequestSchema, async (request) => ({
|
||||
protocolVersion: request.params.protocolVersion,
|
||||
capabilities: {},
|
||||
serverInfo: {
|
||||
name: 'test-server',
|
||||
version: '1.0.0',
|
||||
}
|
||||
}));
|
||||
|
||||
server.setRequestHandler(ListToolsRequestSchema, async () => ({
|
||||
tools: [
|
||||
{
|
||||
name: 'test-tool',
|
||||
description: 'A test tool',
|
||||
inputSchema: {
|
||||
type: 'object',
|
||||
properties: {},
|
||||
},
|
||||
outputSchema: {
|
||||
type: 'object',
|
||||
properties: {
|
||||
result: { type: 'string' },
|
||||
count: { type: 'number' },
|
||||
},
|
||||
required: ['result', 'count'],
|
||||
additionalProperties: false,
|
||||
},
|
||||
},
|
||||
],
|
||||
}));
|
||||
|
||||
server.setRequestHandler(CallToolRequestSchema, async (request) => {
|
||||
if (request.params.name === 'test-tool') {
|
||||
return {
|
||||
structuredContent: JSON.stringify({ result: 'success', count: 42 }),
|
||||
};
|
||||
}
|
||||
throw new Error('Unknown tool');
|
||||
});
|
||||
|
||||
const client = new Client({
|
||||
name: 'test-client',
|
||||
version: '1.0.0',
|
||||
});
|
||||
|
||||
const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair();
|
||||
|
||||
await Promise.all([
|
||||
client.connect(clientTransport),
|
||||
server.connect(serverTransport),
|
||||
]);
|
||||
|
||||
// List tools to cache the schemas
|
||||
await client.listTools();
|
||||
|
||||
// Call the tool - should validate successfully
|
||||
const result = await client.callTool({ name: 'test-tool' });
|
||||
expect(result.structuredContent).toBe('{"result":"success","count":42}');
|
||||
});
|
||||
|
||||
test('should throw error when structuredContent does not match schema', async () => {
|
||||
const server = new Server({
|
||||
name: 'test-server',
|
||||
version: '1.0.0',
|
||||
}, {
|
||||
capabilities: {
|
||||
tools: {},
|
||||
},
|
||||
});
|
||||
|
||||
// Set up server handlers
|
||||
server.setRequestHandler(InitializeRequestSchema, async (request) => ({
|
||||
protocolVersion: request.params.protocolVersion,
|
||||
capabilities: {},
|
||||
serverInfo: {
|
||||
name: 'test-server',
|
||||
version: '1.0.0',
|
||||
}
|
||||
}));
|
||||
|
||||
server.setRequestHandler(ListToolsRequestSchema, async () => ({
|
||||
tools: [
|
||||
{
|
||||
name: 'test-tool',
|
||||
description: 'A test tool',
|
||||
inputSchema: {
|
||||
type: 'object',
|
||||
properties: {},
|
||||
},
|
||||
outputSchema: {
|
||||
type: 'object',
|
||||
properties: {
|
||||
result: { type: 'string' },
|
||||
count: { type: 'number' },
|
||||
},
|
||||
required: ['result', 'count'],
|
||||
additionalProperties: false,
|
||||
},
|
||||
},
|
||||
],
|
||||
}));
|
||||
|
||||
server.setRequestHandler(CallToolRequestSchema, async (request) => {
|
||||
if (request.params.name === 'test-tool') {
|
||||
// Return invalid structured content (count is string instead of number)
|
||||
return {
|
||||
structuredContent: JSON.stringify({ result: 'success', count: 'not a number' }),
|
||||
};
|
||||
}
|
||||
throw new Error('Unknown tool');
|
||||
});
|
||||
|
||||
const client = new Client({
|
||||
name: 'test-client',
|
||||
version: '1.0.0',
|
||||
});
|
||||
|
||||
const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair();
|
||||
|
||||
await Promise.all([
|
||||
client.connect(clientTransport),
|
||||
server.connect(serverTransport),
|
||||
]);
|
||||
|
||||
// List tools to cache the schemas
|
||||
await client.listTools();
|
||||
|
||||
// Call the tool - should throw validation error
|
||||
await expect(client.callTool({ name: 'test-tool' })).rejects.toThrow(
|
||||
/Structured content does not match the tool's output schema/
|
||||
);
|
||||
});
|
||||
|
||||
test('should throw error when tool with outputSchema returns no structuredContent', async () => {
|
||||
const server = new Server({
|
||||
name: 'test-server',
|
||||
version: '1.0.0',
|
||||
}, {
|
||||
capabilities: {
|
||||
tools: {},
|
||||
},
|
||||
});
|
||||
|
||||
// Set up server handlers
|
||||
server.setRequestHandler(InitializeRequestSchema, async (request) => ({
|
||||
protocolVersion: request.params.protocolVersion,
|
||||
capabilities: {},
|
||||
serverInfo: {
|
||||
name: 'test-server',
|
||||
version: '1.0.0',
|
||||
}
|
||||
}));
|
||||
|
||||
server.setRequestHandler(ListToolsRequestSchema, async () => ({
|
||||
tools: [
|
||||
{
|
||||
name: 'test-tool',
|
||||
description: 'A test tool',
|
||||
inputSchema: {
|
||||
type: 'object',
|
||||
properties: {},
|
||||
},
|
||||
outputSchema: {
|
||||
type: 'object',
|
||||
properties: {
|
||||
result: { type: 'string' },
|
||||
},
|
||||
required: ['result'],
|
||||
},
|
||||
},
|
||||
],
|
||||
}));
|
||||
|
||||
server.setRequestHandler(CallToolRequestSchema, async (request) => {
|
||||
if (request.params.name === 'test-tool') {
|
||||
// Return content instead of structuredContent
|
||||
return {
|
||||
content: [{ type: 'text', text: 'This should be structured content' }],
|
||||
};
|
||||
}
|
||||
throw new Error('Unknown tool');
|
||||
});
|
||||
|
||||
const client = new Client({
|
||||
name: 'test-client',
|
||||
version: '1.0.0',
|
||||
});
|
||||
|
||||
const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair();
|
||||
|
||||
await Promise.all([
|
||||
client.connect(clientTransport),
|
||||
server.connect(serverTransport),
|
||||
]);
|
||||
|
||||
// List tools to cache the schemas
|
||||
await client.listTools();
|
||||
|
||||
// Call the tool - should throw error
|
||||
await expect(client.callTool({ name: 'test-tool' })).rejects.toThrow(
|
||||
/Tool test-tool has an output schema but did not return structured content/
|
||||
);
|
||||
});
|
||||
|
||||
test('should handle tools without outputSchema normally', async () => {
|
||||
const server = new Server({
|
||||
name: 'test-server',
|
||||
version: '1.0.0',
|
||||
}, {
|
||||
capabilities: {
|
||||
tools: {},
|
||||
},
|
||||
});
|
||||
|
||||
// Set up server handlers
|
||||
server.setRequestHandler(InitializeRequestSchema, async (request) => ({
|
||||
protocolVersion: request.params.protocolVersion,
|
||||
capabilities: {},
|
||||
serverInfo: {
|
||||
name: 'test-server',
|
||||
version: '1.0.0',
|
||||
}
|
||||
}));
|
||||
|
||||
server.setRequestHandler(ListToolsRequestSchema, async () => ({
|
||||
tools: [
|
||||
{
|
||||
name: 'test-tool',
|
||||
description: 'A test tool',
|
||||
inputSchema: {
|
||||
type: 'object',
|
||||
properties: {},
|
||||
},
|
||||
// No outputSchema
|
||||
},
|
||||
],
|
||||
}));
|
||||
|
||||
server.setRequestHandler(CallToolRequestSchema, async (request) => {
|
||||
if (request.params.name === 'test-tool') {
|
||||
// Return regular content
|
||||
return {
|
||||
content: [{ type: 'text', text: 'Normal response' }],
|
||||
};
|
||||
}
|
||||
throw new Error('Unknown tool');
|
||||
});
|
||||
|
||||
const client = new Client({
|
||||
name: 'test-client',
|
||||
version: '1.0.0',
|
||||
});
|
||||
|
||||
const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair();
|
||||
|
||||
await Promise.all([
|
||||
client.connect(clientTransport),
|
||||
server.connect(serverTransport),
|
||||
]);
|
||||
|
||||
// List tools to cache the schemas
|
||||
await client.listTools();
|
||||
|
||||
// Call the tool - should work normally without validation
|
||||
const result = await client.callTool({ name: 'test-tool' });
|
||||
expect(result.content).toEqual([{ type: 'text', text: 'Normal response' }]);
|
||||
});
|
||||
|
||||
test('should handle complex JSON schema validation', async () => {
|
||||
const server = new Server({
|
||||
name: 'test-server',
|
||||
version: '1.0.0',
|
||||
}, {
|
||||
capabilities: {
|
||||
tools: {},
|
||||
},
|
||||
});
|
||||
|
||||
// Set up server handlers
|
||||
server.setRequestHandler(InitializeRequestSchema, async (request) => ({
|
||||
protocolVersion: request.params.protocolVersion,
|
||||
capabilities: {},
|
||||
serverInfo: {
|
||||
name: 'test-server',
|
||||
version: '1.0.0',
|
||||
}
|
||||
}));
|
||||
|
||||
server.setRequestHandler(ListToolsRequestSchema, async () => ({
|
||||
tools: [
|
||||
{
|
||||
name: 'complex-tool',
|
||||
description: 'A tool with complex schema',
|
||||
inputSchema: {
|
||||
type: 'object',
|
||||
properties: {},
|
||||
},
|
||||
outputSchema: {
|
||||
type: 'object',
|
||||
properties: {
|
||||
name: { type: 'string', minLength: 3 },
|
||||
age: { type: 'integer', minimum: 0, maximum: 120 },
|
||||
active: { type: 'boolean' },
|
||||
tags: {
|
||||
type: 'array',
|
||||
items: { type: 'string' },
|
||||
minItems: 1,
|
||||
},
|
||||
metadata: {
|
||||
type: 'object',
|
||||
properties: {
|
||||
created: { type: 'string' },
|
||||
},
|
||||
required: ['created'],
|
||||
},
|
||||
},
|
||||
required: ['name', 'age', 'active', 'tags', 'metadata'],
|
||||
additionalProperties: false,
|
||||
},
|
||||
},
|
||||
],
|
||||
}));
|
||||
|
||||
server.setRequestHandler(CallToolRequestSchema, async (request) => {
|
||||
if (request.params.name === 'complex-tool') {
|
||||
return {
|
||||
structuredContent: JSON.stringify({
|
||||
name: 'John Doe',
|
||||
age: 30,
|
||||
active: true,
|
||||
tags: ['user', 'admin'],
|
||||
metadata: {
|
||||
created: '2023-01-01T00:00:00Z',
|
||||
},
|
||||
}),
|
||||
};
|
||||
}
|
||||
throw new Error('Unknown tool');
|
||||
});
|
||||
|
||||
const client = new Client({
|
||||
name: 'test-client',
|
||||
version: '1.0.0',
|
||||
});
|
||||
|
||||
const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair();
|
||||
|
||||
await Promise.all([
|
||||
client.connect(clientTransport),
|
||||
server.connect(serverTransport),
|
||||
]);
|
||||
|
||||
// List tools to cache the schemas
|
||||
await client.listTools();
|
||||
|
||||
// Call the tool - should validate successfully
|
||||
const result = await client.callTool({ name: 'complex-tool' });
|
||||
expect(result.structuredContent).toBeDefined();
|
||||
const parsedContent = JSON.parse(result.structuredContent as string);
|
||||
expect(parsedContent.name).toBe('John Doe');
|
||||
expect(parsedContent.age).toBe(30);
|
||||
});
|
||||
|
||||
test('should fail validation with additional properties when not allowed', async () => {
|
||||
const server = new Server({
|
||||
name: 'test-server',
|
||||
version: '1.0.0',
|
||||
}, {
|
||||
capabilities: {
|
||||
tools: {},
|
||||
},
|
||||
});
|
||||
|
||||
// Set up server handlers
|
||||
server.setRequestHandler(InitializeRequestSchema, async (request) => ({
|
||||
protocolVersion: request.params.protocolVersion,
|
||||
capabilities: {},
|
||||
serverInfo: {
|
||||
name: 'test-server',
|
||||
version: '1.0.0',
|
||||
}
|
||||
}));
|
||||
|
||||
server.setRequestHandler(ListToolsRequestSchema, async () => ({
|
||||
tools: [
|
||||
{
|
||||
name: 'strict-tool',
|
||||
description: 'A tool with strict schema',
|
||||
inputSchema: {
|
||||
type: 'object',
|
||||
properties: {},
|
||||
},
|
||||
outputSchema: {
|
||||
type: 'object',
|
||||
properties: {
|
||||
name: { type: 'string' },
|
||||
},
|
||||
required: ['name'],
|
||||
additionalProperties: false,
|
||||
},
|
||||
},
|
||||
],
|
||||
}));
|
||||
|
||||
server.setRequestHandler(CallToolRequestSchema, async (request) => {
|
||||
if (request.params.name === 'strict-tool') {
|
||||
// Return structured content with extra property
|
||||
return {
|
||||
structuredContent: JSON.stringify({
|
||||
name: 'John',
|
||||
extraField: 'not allowed',
|
||||
}),
|
||||
};
|
||||
}
|
||||
throw new Error('Unknown tool');
|
||||
});
|
||||
|
||||
const client = new Client({
|
||||
name: 'test-client',
|
||||
version: '1.0.0',
|
||||
});
|
||||
|
||||
const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair();
|
||||
|
||||
await Promise.all([
|
||||
client.connect(clientTransport),
|
||||
server.connect(serverTransport),
|
||||
]);
|
||||
|
||||
// List tools to cache the schemas
|
||||
await client.listTools();
|
||||
|
||||
// Call the tool - should throw validation error due to additional property
|
||||
await expect(client.callTool({ name: 'strict-tool' })).rejects.toThrow(
|
||||
/Structured content does not match the tool's output schema/
|
||||
);
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -39,7 +39,12 @@ import {
|
|||
SubscribeRequest,
|
||||
SUPPORTED_PROTOCOL_VERSIONS,
|
||||
UnsubscribeRequest,
|
||||
Tool,
|
||||
ErrorCode,
|
||||
McpError,
|
||||
} from "../types.js";
|
||||
import { z } from "zod";
|
||||
import { parseSchema } from "json-schema-to-zod";
|
||||
|
||||
export type ClientOptions = ProtocolOptions & {
|
||||
/**
|
||||
|
|
@ -86,6 +91,8 @@ export class Client<
|
|||
private _serverVersion?: Implementation;
|
||||
private _capabilities: ClientCapabilities;
|
||||
private _instructions?: string;
|
||||
private _cachedTools: Map<string, Tool> = new Map();
|
||||
private _cachedToolOutputSchemas: Map<string, z.ZodTypeAny> = new Map();
|
||||
|
||||
/**
|
||||
* Initializes this client with the given name and version information.
|
||||
|
|
@ -413,22 +420,83 @@ export class Client<
|
|||
| typeof CompatibilityCallToolResultSchema = CallToolResultSchema,
|
||||
options?: RequestOptions,
|
||||
) {
|
||||
return this.request(
|
||||
const result = await this.request(
|
||||
{ method: "tools/call", params },
|
||||
resultSchema,
|
||||
options,
|
||||
);
|
||||
|
||||
// Check if the tool has an outputSchema
|
||||
const outputSchema = this._cachedToolOutputSchemas.get(params.name);
|
||||
if (outputSchema) {
|
||||
// If tool has outputSchema, it MUST return structuredContent
|
||||
if (!result.structuredContent) {
|
||||
throw new McpError(
|
||||
ErrorCode.InvalidRequest,
|
||||
`Tool ${params.name} has an output schema but did not return structured content`
|
||||
);
|
||||
}
|
||||
|
||||
try {
|
||||
// Parse the structured content as JSON
|
||||
const contentData = JSON.parse(result.structuredContent as string);
|
||||
|
||||
// Validate the content against the schema
|
||||
const validationResult = outputSchema.safeParse(contentData);
|
||||
|
||||
if (!validationResult.success) {
|
||||
throw new McpError(
|
||||
ErrorCode.InvalidParams,
|
||||
`Structured content does not match the tool's output schema: ${validationResult.error.message}`
|
||||
);
|
||||
}
|
||||
} catch (error) {
|
||||
if (error instanceof McpError) {
|
||||
throw error;
|
||||
}
|
||||
throw new McpError(
|
||||
ErrorCode.InvalidParams,
|
||||
`Failed to validate structured content: ${error instanceof Error ? error.message : String(error)}`
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
async listTools(
|
||||
params?: ListToolsRequest["params"],
|
||||
options?: RequestOptions,
|
||||
) {
|
||||
return this.request(
|
||||
const result = await this.request(
|
||||
{ method: "tools/list", params },
|
||||
ListToolsResultSchema,
|
||||
options,
|
||||
);
|
||||
|
||||
// Cache the tools and their output schemas for future validation
|
||||
this._cachedTools.clear();
|
||||
this._cachedToolOutputSchemas.clear();
|
||||
|
||||
for (const tool of result.tools) {
|
||||
this._cachedTools.set(tool.name, tool);
|
||||
|
||||
// If the tool has an outputSchema, create and cache the Zod schema
|
||||
if (tool.outputSchema) {
|
||||
try {
|
||||
const zodSchemaCode = parseSchema(tool.outputSchema);
|
||||
// The library returns a string of Zod code, we need to evaluate it
|
||||
// Using Function constructor to safely evaluate the Zod schema
|
||||
const createSchema = new Function('z', `return ${zodSchemaCode}`);
|
||||
const zodSchema = createSchema(z);
|
||||
this._cachedToolOutputSchemas.set(tool.name, zodSchema);
|
||||
} catch (error) {
|
||||
console.warn(`Failed to create Zod schema for tool ${tool.name}: ${error}`);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
async sendRootsListChanged() {
|
||||
|
|
|
|||
Loading…
Reference in a new issue