mcp-framework
Version:
Framework for building Model Context Protocol (MCP) servers in Typescript
402 lines (401 loc) • 14.3 kB
JavaScript
import { z } from 'zod';
/**
* Base class for MCP tools using Zod schemas for input validation and type inference.
*
* Define your tool schema using Zod with descriptions:
* ```typescript
* const schema = z.object({
* message: z.string().describe("The message to process")
* });
*
* class MyTool extends MCPTool {
* name = "my_tool";
* description = "My tool description";
* schema = schema;
*
* async execute(input: McpInput<this>) {
* // input is fully typed from your schema
* return input.message;
* }
* }
* ```
*/
export class MCPTool {
useStringify = true;
/**
* Validates the tool schema. This is called automatically when the tool is registered
* with an MCP server, but can also be called manually for testing.
*/
validate() {
if (this.isZodObjectSchema(this.schema)) {
// Access inputSchema to trigger validation
const _ = this.inputSchema;
}
}
isZodObjectSchema(schema) {
return schema instanceof z.ZodObject;
}
get inputSchema() {
if (this.isZodObjectSchema(this.schema)) {
return this.generateSchemaFromZodObject(this.schema);
}
else {
return this.generateSchemaFromLegacyFormat(this.schema);
}
}
generateSchemaFromZodObject(zodSchema) {
const shape = zodSchema.shape;
const properties = {};
const required = [];
const missingDescriptions = [];
Object.entries(shape).forEach(([key, fieldSchema]) => {
const fieldInfo = this.extractFieldInfo(fieldSchema);
if (!fieldInfo.jsonSchema.description) {
missingDescriptions.push(key);
}
properties[key] = fieldInfo.jsonSchema;
if (!fieldInfo.isOptional) {
required.push(key);
}
});
if (missingDescriptions.length > 0) {
throw new Error(`Missing descriptions for fields in ${this.name}: ${missingDescriptions.join(', ')}. ` +
`All fields must have descriptions when using Zod object schemas. ` +
`Use .describe() on each field, e.g., z.string().describe("Field description")`);
}
return {
type: 'object',
properties,
required,
};
}
extractFieldInfo(schema) {
let currentSchema = schema;
let isOptional = false;
let defaultValue;
let description;
// Extract description before unwrapping
const getDescription = (s) => s._def?.description;
description = getDescription(currentSchema);
// Unwrap modifiers to get to the base type
while (true) {
if (currentSchema instanceof z.ZodOptional) {
isOptional = true;
currentSchema = currentSchema.unwrap();
if (!description)
description = getDescription(currentSchema);
}
else if (currentSchema instanceof z.ZodDefault) {
defaultValue = currentSchema._def.defaultValue();
currentSchema = currentSchema._def.innerType;
if (!description)
description = getDescription(currentSchema);
}
else if (currentSchema instanceof z.ZodNullable) {
isOptional = true;
currentSchema = currentSchema.unwrap();
if (!description)
description = getDescription(currentSchema);
}
else {
break;
}
}
// Build JSON Schema
const jsonSchema = {
type: this.getJsonSchemaTypeFromZod(currentSchema),
};
if (description) {
jsonSchema.description = description;
}
if (defaultValue !== undefined) {
jsonSchema.default = defaultValue;
}
// Handle enums
if (currentSchema instanceof z.ZodEnum) {
jsonSchema.enum = currentSchema._def.values;
}
// Handle arrays
if (currentSchema instanceof z.ZodArray) {
const itemInfo = this.extractFieldInfo(currentSchema._def.type);
jsonSchema.items = itemInfo.jsonSchema;
}
// Handle nested objects
if (currentSchema instanceof z.ZodObject) {
const shape = currentSchema.shape;
const nestedProperties = {};
const nestedRequired = [];
Object.entries(shape).forEach(([key, fieldSchema]) => {
const nestedFieldInfo = this.extractFieldInfo(fieldSchema);
nestedProperties[key] = nestedFieldInfo.jsonSchema;
if (!nestedFieldInfo.isOptional) {
nestedRequired.push(key);
}
});
jsonSchema.properties = nestedProperties;
if (nestedRequired.length > 0) {
jsonSchema.required = nestedRequired;
}
}
// Handle numeric constraints
if (currentSchema instanceof z.ZodNumber) {
const checks = currentSchema._def.checks || [];
checks.forEach((check) => {
switch (check.kind) {
case 'min':
jsonSchema.minimum = check.value;
if (check.inclusive === false) {
jsonSchema.exclusiveMinimum = true;
}
break;
case 'max':
jsonSchema.maximum = check.value;
if (check.inclusive === false) {
jsonSchema.exclusiveMaximum = true;
}
break;
case 'int':
jsonSchema.type = 'integer';
break;
}
});
// Handle positive() which adds a min check of 0 (exclusive)
const hasPositive = checks.some((check) => check.kind === 'min' && check.value === 0 && check.inclusive === false);
if (hasPositive) {
jsonSchema.minimum = 1;
}
}
// Handle string constraints
if (currentSchema instanceof z.ZodString) {
const checks = currentSchema._def.checks || [];
checks.forEach((check) => {
switch (check.kind) {
case 'min':
jsonSchema.minLength = check.value;
break;
case 'max':
jsonSchema.maxLength = check.value;
break;
case 'regex':
jsonSchema.pattern = check.regex.source;
break;
case 'email':
jsonSchema.format = 'email';
break;
case 'url':
jsonSchema.format = 'uri';
break;
case 'uuid':
jsonSchema.format = 'uuid';
break;
}
});
}
return { jsonSchema, isOptional };
}
getJsonSchemaTypeFromZod(zodType) {
if (zodType instanceof z.ZodString)
return 'string';
if (zodType instanceof z.ZodNumber)
return 'number';
if (zodType instanceof z.ZodBoolean)
return 'boolean';
if (zodType instanceof z.ZodArray)
return 'array';
if (zodType instanceof z.ZodObject)
return 'object';
if (zodType instanceof z.ZodEnum)
return 'string';
if (zodType instanceof z.ZodNull)
return 'null';
if (zodType instanceof z.ZodUndefined)
return 'undefined';
if (zodType instanceof z.ZodLiteral) {
const value = zodType._def.value;
return typeof value === 'string'
? 'string'
: typeof value === 'number'
? 'number'
: typeof value === 'boolean'
? 'boolean'
: 'string';
}
return 'string';
}
generateSchemaFromLegacyFormat(schema) {
const properties = {};
const required = [];
Object.entries(schema).forEach(([key, fieldSchema]) => {
// Determine the correct JSON schema type (unwrapping optional if necessary)
const jsonType = this.getJsonSchemaType(fieldSchema.type);
properties[key] = {
type: jsonType,
description: fieldSchema.description,
};
// If the field is not an optional, add it to the required array.
if (!(fieldSchema.type instanceof z.ZodOptional)) {
required.push(key);
}
});
const inputSchema = {
type: 'object',
properties,
};
if (required.length > 0) {
inputSchema.required = required;
}
return inputSchema;
}
get toolDefinition() {
return {
name: this.name,
description: this.description,
inputSchema: this.inputSchema,
};
}
async toolCall(request) {
try {
const args = request.params.arguments || {};
const validatedInput = await this.validateInput(args);
const result = await this.execute(validatedInput);
return this.createSuccessResponse(result);
}
catch (error) {
return this.createErrorResponse(error);
}
}
async validateInput(args) {
if (this.isZodObjectSchema(this.schema)) {
return this.schema.parse(args);
}
else {
const zodSchema = z.object(Object.fromEntries(Object.entries(this.schema).map(([key, schema]) => [
key,
schema.type,
])));
return zodSchema.parse(args);
}
}
getJsonSchemaType(zodType) {
// Unwrap optional types to correctly determine the JSON schema type.
let currentType = zodType;
if (currentType instanceof z.ZodOptional) {
currentType = currentType.unwrap();
}
if (currentType instanceof z.ZodString)
return 'string';
if (currentType instanceof z.ZodNumber)
return 'number';
if (currentType instanceof z.ZodBoolean)
return 'boolean';
if (currentType instanceof z.ZodArray)
return 'array';
if (currentType instanceof z.ZodObject)
return 'object';
return 'string';
}
createSuccessResponse(data) {
if (this.isImageContent(data)) {
return {
content: [data],
};
}
if (Array.isArray(data)) {
const validContent = data.filter((item) => this.isValidContent(item));
if (validContent.length > 0) {
return {
content: validContent,
};
}
}
return {
content: [
{
type: 'text',
text: this.useStringify ? JSON.stringify(data) : String(data),
},
],
};
}
createErrorResponse(error) {
return {
content: [{ type: 'error', text: error.message }],
};
}
isImageContent(data) {
return (typeof data === 'object' &&
data !== null &&
'type' in data &&
data.type === 'image' &&
'data' in data &&
'mimeType' in data &&
typeof data.data === 'string' &&
typeof data.mimeType === 'string');
}
isTextContent(data) {
return (typeof data === 'object' &&
data !== null &&
'type' in data &&
data.type === 'text' &&
'text' in data &&
typeof data.text === 'string');
}
isErrorContent(data) {
return (typeof data === 'object' &&
data !== null &&
'type' in data &&
data.type === 'error' &&
'text' in data &&
typeof data.text === 'string');
}
isValidContent(data) {
return this.isImageContent(data) || this.isTextContent(data) || this.isErrorContent(data);
}
async fetch(url, init) {
const response = await fetch(url, init);
if (!response.ok) {
throw new Error(`HTTP error! status: ${response.status}`);
}
return response.json();
}
}
/**
* Helper function to define tool schemas with required descriptions.
* This ensures all fields have descriptions at build time.
*
* @example
* const schema = defineSchema({
* name: z.string().describe("User's name"),
* age: z.number().describe("User's age")
* });
*/
export function defineSchema(shape) {
// Check descriptions at runtime during development
if (process.env.NODE_ENV !== 'production') {
for (const [key, value] of Object.entries(shape)) {
let schema = value;
let hasDescription = false;
// Check the schema and its wrapped versions for description
while (schema && typeof schema === 'object') {
if ('_def' in schema && schema._def?.description) {
hasDescription = true;
break;
}
// Check wrapped types
if (schema instanceof z.ZodOptional ||
schema instanceof z.ZodDefault ||
schema instanceof z.ZodNullable) {
schema = schema._def.innerType || schema.unwrap();
}
else {
break;
}
}
if (!hasDescription) {
throw new Error(`Field '${key}' is missing a description. Use .describe() to add one.\n` +
`Example: ${key}: z.string().describe("Description for ${key}")`);
}
}
}
return z.object(shape);
}