From 8c37d4ba3b3ad933d2ba898b9c5fc18eeb2dfce1 Mon Sep 17 00:00:00 2001
From: Henry <hzj94@hotmail.com>
Date: Fri, 18 Jul 2025 13:32:50 +0100
Subject: [PATCH] feat: Implement caching for MCP toolkit in CachePool

- Added methods to CachePool for managing MCP toolkit cache.
- Integrated caching logic in CustomMCP to store and retrieve toolkit data based on workspaceId and configuration.
- Updated node service to pass cachePool to CustomMCP for enhanced performance.
---
 .../nodes/tools/MCP/CustomMCP/CustomMCP.ts    | 23 ++++++++++++++++++
 packages/server/src/CachePool.ts              | 24 +++++++++++++++++++
 packages/server/src/services/nodes/index.ts   |  3 ++-
 3 files changed, 49 insertions(+), 1 deletion(-)

diff --git a/packages/components/nodes/tools/MCP/CustomMCP/CustomMCP.ts b/packages/components/nodes/tools/MCP/CustomMCP/CustomMCP.ts
index d81c5d778..b10ece772 100644
--- a/packages/components/nodes/tools/MCP/CustomMCP/CustomMCP.ts
+++ b/packages/components/nodes/tools/MCP/CustomMCP/CustomMCP.ts
@@ -3,6 +3,7 @@ import { ICommonObject, IDatabaseEntity, INode, INodeData, INodeOptionsValue, IN
 import { MCPToolkit } from '../core'
 import { getVars, prepareSandboxVars } from '../../../../src/utils'
 import { DataSource } from 'typeorm'
+import hash from 'object-hash'
 
 const mcpServerConfig = `{
     "command": "npx",
@@ -140,6 +141,24 @@ class Custom_MCP implements INode {
             sandbox['$vars'] = prepareSandboxVars(variables)
         }
 
+        const workspaceId = options?.searchOptions?.workspaceId?._value || options?.workspaceId
+
+        let canonicalConfig
+        try {
+            canonicalConfig = JSON.parse(mcpServerConfig)
+        } catch (e) {
+            canonicalConfig = mcpServerConfig
+        }
+
+        const cacheKey = hash({ workspaceId, canonicalConfig, sandbox })
+
+        if (options.cachePool) {
+            const cachedResult = await options.cachePool.getMCPCache(cacheKey)
+            if (cachedResult) {
+                return cachedResult.tools
+            }
+        }
+
         try {
             let serverParams
             if (typeof mcpServerConfig === 'object') {
@@ -162,6 +181,10 @@ class Custom_MCP implements INode {
 
             const tools = toolkit.tools ?? []
 
+            if (options.cachePool) {
+                await options.cachePool.addMCPCache(cacheKey, { toolkit, tools })
+            }
+
             return tools as Tool[]
         } catch (error) {
             throw new Error(`Invalid MCP Server Config: ${error}`)
diff --git a/packages/server/src/CachePool.ts b/packages/server/src/CachePool.ts
index e978d89de..e8e7f600e 100644
--- a/packages/server/src/CachePool.ts
+++ b/packages/server/src/CachePool.ts
@@ -8,6 +8,7 @@ export class CachePool {
     private redisClient: Redis | null = null
     activeLLMCache: IActiveCache = {}
     activeEmbeddingCache: IActiveCache = {}
+    activeMCPCache: { [key: string]: any } = {}
 
     constructor() {
         if (process.env.MODE === MODE.QUEUE) {
@@ -73,6 +74,29 @@ export class CachePool {
         }
     }
 
+    /**
+     * Add to the mcp toolkit cache pool
+     * @param {string} cacheKey
+     * @param {any} value
+     */
+    async addMCPCache(cacheKey: string, value: any) {
+        // Only add to cache for non-queue mode, because we are storing the toolkit instances in memory, and we can't store them in redis
+        if (process.env.MODE !== MODE.QUEUE) {
+            this.activeMCPCache[`mcpCache:${cacheKey}`] = value
+        }
+    }
+
+    /**
+     * Get item from mcp toolkit cache pool
+     * @param {string} cacheKey
+     */
+    async getMCPCache(cacheKey: string): Promise<any | undefined> {
+        if (process.env.MODE !== MODE.QUEUE) {
+            return this.activeMCPCache[`mcpCache:${cacheKey}`]
+        }
+        return undefined
+    }
+
     /**
      * Get item from llm cache pool
      * @param {string} chatflowid
diff --git a/packages/server/src/services/nodes/index.ts b/packages/server/src/services/nodes/index.ts
index 9a64cb55b..fbe08368a 100644
--- a/packages/server/src/services/nodes/index.ts
+++ b/packages/server/src/services/nodes/index.ts
@@ -103,7 +103,8 @@ const getSingleNodeAsyncOptions = async (nodeName: string, requestBody: any): Pr
                     componentNodes: appServer.nodesPool.componentNodes,
                     previousNodes: requestBody.previousNodes,
                     currentNode: requestBody.currentNode,
-                    searchOptions: requestBody.searchOptions
+                    searchOptions: requestBody.searchOptions,
+                    cachePool: appServer.cachePool
                 })
 
                 return dbResponse