diff --git a/packages/components/nodes/tools/MCP/core.test.ts b/packages/components/nodes/tools/MCP/core.test.ts index 3fcf44d31af..2c2908b5df0 100644 --- a/packages/components/nodes/tools/MCP/core.test.ts +++ b/packages/components/nodes/tools/MCP/core.test.ts @@ -3,10 +3,211 @@ import { validateCommandInjection, validateArgsForLocalFileAccess, validateEnvironmentVariables, + MCPTool, + mcpInputSchemaToZodObject, validateMCPServerConfig } from './core' describe('MCP Security Validations', () => { + describe('mcpInputSchemaToZodObject', () => { + it('should preserve MCP input schema descriptions and required fields', () => { + const schema = mcpInputSchemaToZodObject({ + type: 'object', + properties: { + location: { + type: 'string', + description: 'The city and country to get weather for.' + }, + days: { + type: 'integer', + description: 'Number of forecast days.' + }, + unit: { + type: 'string', + enum: ['celsius', 'fahrenheit'], + description: 'Temperature unit.' + } + }, + required: ['location'] + }) + + expect(schema.safeParse({}).success).toBe(false) + expect(schema.safeParse({ location: 'Dhaka, Bangladesh', unit: 'kelvin' }).success).toBe(false) + expect(schema.safeParse({ location: 'Dhaka, Bangladesh', days: 3, unit: 'celsius' }).success).toBe(true) + + const shape = schema.shape + expect(shape.location.description).toBe('The city and country to get weather for.') + expect(shape.days.description).toBe('Number of forecast days.') + expect(shape.unit.description).toBe('Temperature unit.') + expect(shape.location.safeParse(undefined).success).toBe(false) + expect(shape.days.safeParse(undefined).success).toBe(true) + }) + + it('should preserve nested object and array parameter schemas', () => { + const schema = mcpInputSchemaToZodObject({ + type: 'object', + properties: { + filters: { + type: 'object', + description: 'Search filters.', + properties: { + tags: { + type: 'array', + description: 'Tags to match.', + items: { + type: 'string', + description: 'A tag value.' + } + } + }, + required: ['tags'] + } + }, + required: ['filters'] + }) + + expect(schema.safeParse({ filters: { tags: ['docs', 'api'] } }).success).toBe(true) + expect(schema.safeParse({ filters: {} }).success).toBe(false) + expect(schema.shape.filters.description).toBe('Search filters.') + expect(schema.shape.filters.shape.tags.description).toBe('Tags to match.') + }) + + it('should preserve non-string enum values', () => { + const schema = mcpInputSchemaToZodObject({ + type: 'object', + properties: { + priority: { + enum: [1, 2, 'urgent', true, null], + description: 'Priority marker.' + } + }, + required: ['priority'] + }) + + expect(schema.safeParse({ priority: 1 }).success).toBe(true) + expect(schema.safeParse({ priority: true }).success).toBe(true) + expect(schema.safeParse({ priority: null }).success).toBe(true) + expect(schema.safeParse({ priority: 'urgent' }).success).toBe(true) + expect(schema.safeParse({ priority: 'low' }).success).toBe(false) + expect(schema.shape.priority.description).toBe('Priority marker.') + }) + + it('should support array-based JSON schema types', () => { + const schema = mcpInputSchemaToZodObject({ + type: 'object', + properties: { + query: { + type: ['string', 'null'], + description: 'Optional query text.' + } + }, + required: ['query'] + }) + + expect(schema.safeParse({ query: 'docs' }).success).toBe(true) + expect(schema.safeParse({ query: null }).success).toBe(true) + expect(schema.safeParse({ query: 123 }).success).toBe(false) + expect(schema.safeParse({}).success).toBe(false) + expect(schema.shape.query.description).toBe('Optional query text.') + }) + + it('should infer object and array schemas from structural fields', () => { + const schema = mcpInputSchemaToZodObject({ + type: 'object', + properties: { + filters: { + description: 'Search filters.', + properties: { + tags: { + description: 'Tags to match.', + items: { + type: 'string' + } + } + }, + required: ['tags'] + } + }, + required: ['filters'] + }) + + expect(schema.safeParse({ filters: { tags: ['docs', 'api'] } }).success).toBe(true) + expect(schema.safeParse({ filters: { tags: 'docs' } }).success).toBe(false) + expect(schema.safeParse({ filters: {} }).success).toBe(false) + expect(schema.shape.filters.description).toBe('Search filters.') + expect(schema.shape.filters.shape.tags.description).toBe('Tags to match.') + }) + + it('should preserve default values', () => { + const schema = mcpInputSchemaToZodObject({ + type: 'object', + properties: { + limit: { + type: 'integer', + default: 10, + description: 'Maximum number of results.' + } + } + }) + + const result = schema.safeParse({}) + + expect(result.success).toBe(true) + expect(result.data.limit).toBe(10) + expect(schema.safeParse({ limit: 5 }).success).toBe(true) + expect(schema.safeParse({ limit: 1.5 }).success).toBe(false) + expect(schema.shape.limit.description).toBe('Maximum number of results.') + }) + + it('should throw for invalid or unsupported JSON schemas', () => { + expect(() => { + mcpInputSchemaToZodObject(null) + }).toThrow('Invalid MCP input schema') + + expect(() => { + mcpInputSchemaToZodObject({ + type: 'object', + properties: { + bad: null + } + }) + }).toThrow('Invalid schema definition for property: bad') + + expect(() => { + mcpInputSchemaToZodObject({ + type: 'object', + properties: { + bad: { + type: 'symbol' + } + } + }) + }).toThrow('Unsupported schema type: symbol for property: bad') + }) + + it('should create LangChain MCP tools with the original MCP description and converted schema', async () => { + const mcpTool = await MCPTool({ + toolkit: {} as any, + name: 'buscar_previsao_tempo', + description: 'Retorna a previsão do tempo para os próximos dias.', + argsSchema: { + type: 'object', + properties: { + local: { + type: 'string', + description: 'Local para consultar a previsão.' + } + }, + required: ['local'] + } + }) + + expect(mcpTool.description).toBe('Retorna a previsão do tempo para os próximos dias.') + expect((mcpTool.schema as any).shape.local.description).toBe('Local para consultar a previsão.') + expect((mcpTool.schema as any).safeParse({}).success).toBe(false) + }) + }) + describe('validateCommandFlags', () => { describe('npx command', () => { it('should block -c flag', () => { diff --git a/packages/components/nodes/tools/MCP/core.ts b/packages/components/nodes/tools/MCP/core.ts index 82cca4dfcb7..d002376fc86 100644 --- a/packages/components/nodes/tools/MCP/core.ts +++ b/packages/components/nodes/tools/MCP/core.ts @@ -4,6 +4,7 @@ import { SSEClientTransport } from '@modelcontextprotocol/sdk/client/sse.js' import { StdioClientTransport, StdioServerParameters } from '@modelcontextprotocol/sdk/client/stdio.js' import { StreamableHTTPClientTransport } from '@modelcontextprotocol/sdk/client/streamableHttp.js' import { CallToolRequest, CallToolResultSchema, ListToolsResult, ListToolsResultSchema } from '@modelcontextprotocol/sdk/types.js' +import { z, type ZodTypeAny } from 'zod' import { checkDenyList, secureFetch } from '../../../src/httpSecurity' export class MCPToolkit extends BaseToolkit { @@ -173,11 +174,175 @@ export async function MCPTool({ { name: name, description: description, - schema: argsSchema + schema: mcpInputSchemaToZodObject(argsSchema) } ) } +const isZodSchema = (schema: unknown): schema is ZodTypeAny => { + return typeof schema === 'object' && schema !== null && typeof (schema as ZodTypeAny).safeParse === 'function' +} + +const applyDescription = (schema: T, description?: string): T => { + return description ? (schema.describe(description) as T) : schema +} + +const applyRequiredAndDefault = (schema: ZodTypeAny, required: boolean, defaultValue: unknown): ZodTypeAny => { + if (defaultValue !== undefined) { + return schema.default(defaultValue) + } + + return required ? schema : schema.optional() +} + +const createUnionSchema = (variants: ZodTypeAny[], propertyName: string): ZodTypeAny => { + if (variants.length === 0) { + throw new Error(`No schema variants defined for property: ${propertyName}`) + } + + if (variants.length === 1) { + return variants[0] + } + + return z.union(variants as [ZodTypeAny, ZodTypeAny, ...ZodTypeAny[]]) +} + +const isJsonLiteral = (value: unknown): value is string | number | boolean | null => { + return value === null || ['string', 'number', 'boolean'].includes(typeof value) +} + +const getSchemaType = (schema: any, propertyName: string): string | string[] => { + if (schema.type === undefined) { + if (schema.properties && typeof schema.properties === 'object' && !Array.isArray(schema.properties)) { + return 'object' + } + if (schema.items !== undefined) { + return 'array' + } + + throw new Error(`Schema type could not be inferred for property: ${propertyName}`) + } + + if (Array.isArray(schema.type)) { + if (schema.type.length === 0 || !schema.type.every((type: unknown) => typeof type === 'string')) { + throw new Error(`Invalid schema type for property: ${propertyName}`) + } + + return schema.type + } + + if (typeof schema.type !== 'string') { + throw new Error(`Invalid schema type for property: ${propertyName}`) + } + + return schema.type +} + +const buildJsonSchemaPropertyZod = (schema: any, propertyName: string): ZodTypeAny => { + if (!schema || typeof schema !== 'object' || Array.isArray(schema)) { + throw new Error(`Invalid schema definition for property: ${propertyName}`) + } + + if (Array.isArray(schema.enum) && schema.enum.length > 0) { + const allStrings = schema.enum.every((value: unknown) => typeof value === 'string') + if (allStrings) { + return z.enum(schema.enum as [string, ...string[]]) + } + + if (!schema.enum.every(isJsonLiteral)) { + throw new Error(`Unsupported enum value for property: ${propertyName}`) + } + + return createUnionSchema( + schema.enum.map((value: string | number | boolean | null) => z.literal(value)), + propertyName + ) + } + + if (Array.isArray(schema.oneOf) || Array.isArray(schema.anyOf)) { + return createUnionSchema( + (schema.oneOf ?? schema.anyOf).map((subSchema: any) => buildJsonSchemaPropertyZod(subSchema, propertyName)), + propertyName + ) + } + + const type = getSchemaType(schema, propertyName) + if (Array.isArray(type)) { + return createUnionSchema( + type.map((schemaType) => buildJsonSchemaPropertyZod({ ...schema, type: schemaType }, propertyName)), + propertyName + ) + } + + let zodSchema: ZodTypeAny + switch (type) { + case 'string': + zodSchema = z.string() + break + case 'number': + zodSchema = z.number() + break + case 'integer': + zodSchema = z.number().int() + break + case 'boolean': + zodSchema = z.boolean() + break + case 'array': + zodSchema = z.array(schema.items === undefined ? z.any() : buildJsonSchemaPropertyZod(schema.items, propertyName)) + break + case 'object': + zodSchema = mcpInputSchemaToZodObject(schema) + break + case 'null': + zodSchema = z.null() + break + default: + throw new Error(`Unsupported schema type: ${type} for property: ${propertyName}`) + } + + return zodSchema +} + +const jsonSchemaPropertyToZod = (schema: any, propertyName: string, required: boolean): ZodTypeAny => { + const zodSchema = buildJsonSchemaPropertyZod(schema, propertyName) + const description = schema.description ?? schema.title ?? propertyName + + return applyDescription(applyRequiredAndDefault(zodSchema, required, schema.default), description) +} + +export const mcpInputSchemaToZodObject = (inputSchema: any): any => { + if (isZodSchema(inputSchema)) { + return inputSchema + } + + if (!inputSchema || typeof inputSchema !== 'object' || Array.isArray(inputSchema)) { + throw new Error('Invalid MCP input schema') + } + + const type = inputSchema.type ?? 'object' + if (type !== 'object') { + throw new Error(`Unsupported MCP input schema type: ${type}`) + } + + const properties = inputSchema?.properties + if (properties === undefined) { + return z.object({}) + } + + if (typeof properties !== 'object' || Array.isArray(properties)) { + throw new Error('Invalid MCP input schema properties') + } + + const requiredProperties = new Set(Array.isArray(inputSchema.required) ? inputSchema.required : []) + const zodShape: Record = {} + for (const [propertyName, propertySchema] of Object.entries(properties)) { + zodShape[propertyName] = jsonSchemaPropertyToZod(propertySchema, propertyName, requiredProperties.has(propertyName)) + } + + return z.object(zodShape) +} + export const validateArgsForLocalFileAccess = (args: string[]): void => { const dangerousPatterns = [ // Absolute paths