From 5930f1119c655bcf8d2200ae827a1f5b9fec81d0 Mon Sep 17 00:00:00 2001
From: Henry Heng <henryheng@flowiseai.com>
Date: Fri, 1 Aug 2025 23:48:21 +0100
Subject: [PATCH] Chore/Add CUSTOM_MCP_PROTOCOL (#5003)

add CUSTOM_MCP_PROTOCOL
---
 .../nodes/tools/MCP/CustomMCP/CustomMCP.ts    | 10 +++----
 packages/components/nodes/tools/MCP/core.ts   | 26 -------------------
 2 files changed, 4 insertions(+), 32 deletions(-)

diff --git a/packages/components/nodes/tools/MCP/CustomMCP/CustomMCP.ts b/packages/components/nodes/tools/MCP/CustomMCP/CustomMCP.ts
index a3e9e3a80..3fa177bfb 100644
--- a/packages/components/nodes/tools/MCP/CustomMCP/CustomMCP.ts
+++ b/packages/components/nodes/tools/MCP/CustomMCP/CustomMCP.ts
@@ -1,6 +1,6 @@
 import { Tool } from '@langchain/core/tools'
 import { ICommonObject, IDatabaseEntity, INode, INodeData, INodeOptionsValue, INodeParams } from '../../../../src/Interface'
-import { MCPToolkit, validateMCPServerSecurity } from '../core'
+import { MCPToolkit } from '../core'
 import { getVars, prepareSandboxVars } from '../../../../src/utils'
 import { DataSource } from 'typeorm'
 import hash from 'object-hash'
@@ -173,13 +173,11 @@ class Custom_MCP implements INode {
                 serverParams = JSON.parse(serverParamsString)
             }
 
-            if (process.env.CUSTOM_MCP_SECURITY_CHECK === 'true') {
-                validateMCPServerSecurity(serverParams)
-            }
-
             // Compatible with stdio and SSE
             let toolkit: MCPToolkit
-            if (serverParams?.command === undefined) {
+            if (process.env.CUSTOM_MCP_PROTOCOL === 'sse') {
+                toolkit = new MCPToolkit(serverParams, 'sse')
+            } else if (serverParams?.command === undefined) {
                 toolkit = new MCPToolkit(serverParams, 'sse')
             } else {
                 toolkit = new MCPToolkit(serverParams, 'stdio')
diff --git a/packages/components/nodes/tools/MCP/core.ts b/packages/components/nodes/tools/MCP/core.ts
index a34d44efa..b2c9e63f6 100644
--- a/packages/components/nodes/tools/MCP/core.ts
+++ b/packages/components/nodes/tools/MCP/core.ts
@@ -174,32 +174,6 @@ function createSchemaModel(
     return z.object(schemaProperties)
 }
 
-/**
- * TODO: To be removed and only allow Remote MCP for Cloud
- * Validates MCP server configuration to only allow whitelisted commands
- */
-export function validateMCPServerSecurity(serverParams: Record<string, any>): void {
-    // Whitelist of allowed commands - only these are permitted
-    const allowedCommands = ['npx', 'node']
-
-    if (serverParams.command) {
-        const cmd = serverParams.command.toLowerCase()
-        const baseCmd = cmd
-
-        if (!allowedCommands.includes(baseCmd)) {
-            throw new Error(`Only allowed: ${allowedCommands.join(', ')}`)
-        }
-    }
-
-    if (serverParams.env) {
-        for (const [key, value] of Object.entries(serverParams.env)) {
-            if (typeof value === 'string' && (value.includes('$(') || value.includes('`'))) {
-                throw new Error(`Environment variable "${key}" contains command substitution: "${value}"`)
-            }
-        }
-    }
-}
-
 export const validateArgsForLocalFileAccess = (args: string[]): void => {
     const dangerousPatterns = [
         // Absolute paths