fix(mcp): preserve tool input schema types and required fields (#6248)

* fix(mcp): preserve tool input schema types and required fields

* Update OpenAIAssistant.ts

* replace zod-to-json-schema with toolSchemaToJsonSchema for schema normalization

* enhance isZodSchema validation and handle null schemas in toolSchemaToJsonSchema

Author: HenryHengZJ

packages/components/nodes/agentflow/Agent/Agent.ts

@@ -19,7 +19,7 @@ import { ILLMMessage, IResponseMetadata } from '../Interface.Agentflow'
 import { Tool } from '@langchain/core/tools'
 import { ARTIFACTS_PREFIX, SOURCE_DOCUMENTS_PREFIX, TOOL_ARGS_PREFIX } from '../../../src/agents'
 import { flatten } from 'lodash'
-import zodToJsonSchema from 'zod-to-json-schema'
+import { toolSchemaToJsonSchema, type ToolJsonSchema } from '../../../src/utils'
 import { getErrorMessage } from '../../../src/error'
 import { DataSource } from 'typeorm'
 import { randomBytes } from 'crypto'
@@ -68,7 +68,7 @@ interface IKnowledgeBaseVSEmbeddings {
 interface ISimpliefiedTool {
     name: string
     description: string
-    schema: any
+    schema: ToolJsonSchema
     toolNode: {
         label: string
         name: string
@@ -744,10 +744,7 @@ class Agent_Agentflow implements INode {
                 }
                 const componentNode = options.componentNodes[agentSelectedTool]
 
-                const jsonSchema = zodToJsonSchema(tool.schema as any)
-                if (jsonSchema.$schema) {
-                    delete jsonSchema.$schema
-                }
+                const jsonSchema = toolSchemaToJsonSchema(tool.schema)
 
                 return {
                     name: tool.name,
@@ -801,10 +798,7 @@ class Agent_Agentflow implements INode {
 
                     toolsInstance.push(retrieverToolInstance as Tool)
 
-                    const jsonSchema = zodToJsonSchema(retrieverToolInstance.schema)
-                    if (jsonSchema.$schema) {
-                        delete jsonSchema.$schema
-                    }
+                    const jsonSchema = toolSchemaToJsonSchema(retrieverToolInstance.schema)
                     const componentNode = options.componentNodes['retrieverTool']
 
                     availableTools.push({
@@ -876,10 +870,7 @@ class Agent_Agentflow implements INode {
 
                     toolsInstance.push(retrieverToolInstance as Tool)
 
-                    const jsonSchema = zodToJsonSchema(retrieverToolInstance.schema)
-                    if (jsonSchema.$schema) {
-                        delete jsonSchema.$schema
-                    }
+                    const jsonSchema = toolSchemaToJsonSchema(retrieverToolInstance.schema)
                     const componentNode = options.componentNodes['retrieverTool']
 
                     availableTools.push({

packages/components/nodes/agentflow/Tool/Tool.ts

@@ -3,7 +3,7 @@ import { updateFlowState } from '../utils'
 import { processTemplateVariables } from '../../../src/utils'
 import { Tool } from '@langchain/core/tools'
 import { ARTIFACTS_PREFIX, TOOL_ARGS_PREFIX } from '../../../src/agents'
-import zodToJsonSchema from 'zod-to-json-schema'
+import { toolSchemaToJsonSchema } from '../../../src/utils'
 
 interface IToolInputArgs {
     inputArgName: string
@@ -153,19 +153,16 @@ class Tool_Agentflow implements INode {
                     // Combine schemas from all tools in the array
                     const allProperties = toolInstance.reduce((acc, tool) => {
                         if (tool?.schema) {
-                            const schema: Record<string, any> = zodToJsonSchema(tool.schema)
+                            const schema = toolSchemaToJsonSchema(tool.schema) as { properties?: ICommonObject }
                             return { ...acc, ...(schema.properties || {}) }
                         }
                         return acc
                     }, {})
                     toolInputArgs = { properties: allProperties }
+                } else if (!toolInstance.schema) {
+                    toolInputArgs = {}
                 } else {
-                    // Handle single tool instance
-                    toolInputArgs = toolInstance.schema ? zodToJsonSchema(toolInstance.schema as any) : {}
-                }
-
-                if (toolInputArgs && Object.keys(toolInputArgs).length > 0) {
-                    delete toolInputArgs.$schema
+                    toolInputArgs = toolSchemaToJsonSchema(toolInstance.schema) as ICommonObject
                 }
 
                 return Object.keys(toolInputArgs.properties || {}).map((item) => ({

packages/components/nodes/agents/OpenAIAssistant/OpenAIAssistant.ts

@@ -13,7 +13,7 @@ import { DataSource } from 'typeorm'
 import { getCredentialData, getCredentialParam } from '../../../src/utils'
 import fetch from 'node-fetch'
 import { flatten, uniqWith, isEqual } from 'lodash'
-import { zodToJsonSchema } from 'zod-to-json-schema'
+import { toolSchemaToJsonSchema } from '../../../src/utils'
 import { AnalyticHandler } from '../../../src/handler'
 import { Moderation, checkInputs, streamResponse } from '../../moderation/Moderation'
 import { formatResponse } from '../../outputparsers/OutputParserHelpers'
@@ -1163,7 +1163,7 @@ interface JSONSchema {
 }
 
 const formatToOpenAIAssistantTool = (tool: any): OpenAI.Beta.FunctionTool => {
-    const parameters = zodToJsonSchema(tool.schema) as JSONSchema
+    const parameters = toolSchemaToJsonSchema(tool.schema) as JSONSchema
 
     // For strict tools, we need to:
     // 1. Set additionalProperties to false

packages/components/nodes/tools/MCP/Pipedream/PipedreamMCP.ts

@@ -4,7 +4,6 @@ import { getCredentialData, getCredentialParam, getVars, prepareSandboxVars } fr
 import { DataSource } from 'typeorm'
 import { MCPToolkit } from '../core'
 import axios from 'axios'
-import { z, ZodTypeAny } from 'zod'
 import { CallToolResultSchema } from '@modelcontextprotocol/sdk/types.js'
 import type { CallToolRequest, CallToolResult, TextContent, Tool as McpTool } from '@modelcontextprotocol/sdk/types.js'
 
@@ -40,17 +39,6 @@ Once the user has connected their account, retry the original request.`
     return text
 }
 
-function createSchemaModel(inputSchema: { type: string; properties?: Record<string, any> }): z.ZodObject<any> {
-    if (inputSchema.type !== 'object' || !inputSchema.properties) {
-        throw new Error('Invalid schema type or missing properties')
-    }
-    const schemaProperties = Object.entries(inputSchema.properties).reduce((acc, [key]) => {
-        acc[key] = z.any()
-        return acc
-    }, {} as Record<string, ZodTypeAny>)
-    return z.object(schemaProperties)
-}
-
 async function createPipedreamTool(toolkit: MCPToolkit, name: string, description: string, argsSchema: any): Promise<Tool> {
     return tool(
         async (input): Promise<string> => {
@@ -345,7 +333,7 @@ class Pipedream_MCP implements INode {
             }
 
             const toolPromises = rawTools.map((t: McpTool) =>
-                createPipedreamTool(toolkit, t.name, t.description || t.name, createSchemaModel(t.inputSchema))
+                createPipedreamTool(toolkit, t.name, t.description || t.name, t.inputSchema ?? { type: 'object', properties: {} })
             )
             const settled = await Promise.allSettled(toolPromises)
             const tools = settled.filter((r): r is PromiseFulfilledResult<Tool> => r.status === 'fulfilled').map((r) => r.value)

packages/components/nodes/tools/MCP/core.ts

@@ -2,7 +2,6 @@ import { CallToolRequest, CallToolResultSchema, ListToolsResult, ListToolsResult
 import { Client } from '@modelcontextprotocol/sdk/client/index.js'
 import { StdioClientTransport, StdioServerParameters } from '@modelcontextprotocol/sdk/client/stdio.js'
 import { BaseToolkit, tool, Tool } from '@langchain/core/tools'
-import { z, type ZodTypeAny } from 'zod/v3'
 import { StreamableHTTPClientTransport } from '@modelcontextprotocol/sdk/client/streamableHttp.js'
 import { SSEClientTransport } from '@modelcontextprotocol/sdk/client/sse.js'
 import { checkDenyList, secureFetch } from '../../../src/httpSecurity'
@@ -124,11 +123,12 @@ export class MCPToolkit extends BaseToolkit {
             if (this.client === null) {
                 throw new Error('Client is not initialized')
             }
+            const argsSchema = tool.inputSchema ?? { type: 'object', properties: {} }
             return await MCPTool({
                 toolkit: this,
                 name: tool.name,
                 description: tool.description || tool.name,
-                argsSchema: createSchemaModel(tool.inputSchema)
+                argsSchema
             })
         })
         const res = await Promise.allSettled(toolsPromises)
@@ -177,24 +177,6 @@ export async function MCPTool({
     )
 }
 
-function createSchemaModel(
-    inputSchema: {
-        type: 'object'
-        properties?: Record<string, unknown>
-    } & { [k: string]: unknown }
-): z.ZodObject<Record<string, ZodTypeAny>> {
-    if (inputSchema.type !== 'object' || !inputSchema.properties) {
-        throw new Error('Invalid schema type or missing properties')
-    }
-
-    const schemaProperties = Object.entries(inputSchema.properties).reduce((acc, [key]) => {
-        acc[key] = z.any()
-        return acc
-    }, {} as Record<string, ZodTypeAny>)
-
-    return z.object(schemaProperties)
-}
-
 export const validateArgsForLocalFileAccess = (args: string[]): void => {
     const dangerousPatterns = [
         // Absolute paths

packages/components/src/utils.ts

@@ -22,6 +22,7 @@ import { NodeVM } from 'vm2'
 import { Sandbox } from '@e2b/code-interpreter'
 import { secureFetch, checkDenyList, secureAxiosRequest } from './httpSecurity'
 import JSON5 from 'json5'
+import zodToJsonSchema, { type JsonSchema7Type } from 'zod-to-json-schema'
 
 export const numberOrExpressionRegex = '^(\\d+\\.?\\d*|{{.*}})$' //return true if string consists only numbers OR expression {{}}
 export const notEmptyRegex = '(.|\\s)*\\S(.|\\s)*' //return true if string is not empty or blank
@@ -2303,3 +2304,39 @@ export const isReasoningModelOpenAI = (name: string): boolean => {
     if (name.includes('gpt-5')) return true
     return false
 }
+
+/**
+ * JSON Schema shape returned by {@link toolSchemaToJsonSchema}, extended with the
+ * optional `$schema` marker that `zod-to-json-schema` emits.
+ */
+export type ToolJsonSchema = JsonSchema7Type & { $schema?: string; [key: string]: unknown }
+
+type ZodToJsonSchemaInput = Parameters<typeof zodToJsonSchema>[0]
+
+/**
+ * Type guard detecting a Zod schema without importing Zod's types directly.
+ *
+ * Using `Parameters<typeof zodToJsonSchema>[0]` keeps the guard compatible with
+ * whichever Zod major version (`^3 || ^4`) TypeScript resolves at the call site.
+ */
+export const isZodSchema = (schema: unknown): schema is ZodToJsonSchemaInput =>
+    typeof schema === 'object' && schema !== null && '_def' in schema && typeof (schema as { parse?: unknown }).parse === 'function'
+
+/**
+ * Normalizes a tool schema into a plain JSON Schema object.
+ *
+ * LangChain tools may expose their `schema` as either a Zod schema (has `_def`)
+ * or an already-plain JSON Schema (e.g. MCP tools). This helper handles both,
+ * deep-clones plain objects to prevent accidental mutation, and strips the
+ * `$schema` marker so the result is safe to embed in LLM tool definitions.
+ */
+export const toolSchemaToJsonSchema = (schema: unknown): ToolJsonSchema => {
+    if (schema == null) return { type: 'object', properties: {} }
+    const jsonSchema: ToolJsonSchema = isZodSchema(schema)
+        ? (zodToJsonSchema(schema) as ToolJsonSchema)
+        : cloneDeep(schema as ToolJsonSchema)
+    if (jsonSchema.$schema) {
+        delete jsonSchema.$schema
+    }
+    return jsonSchema
+}