From db2b92421a7ca89898f0a3fa68ba1c198dec8015 Mon Sep 17 00:00:00 2001
From: kangfenmao <kangfenmao@qq.com>
Date: Tue, 11 Mar 2025 13:52:36 +0800
Subject: [PATCH] refactor(MCP): Simplify IPC handlers and improve server
 management

---
 src/main/ipc.ts                               | 50 ++++++-------------
 .../src/pages/settings/MCPSettings.tsx        |  2 +-
 .../mcpToolUtils.ts => utils/mcp-tools.ts}    |  2 +-
 3 files changed, 16 insertions(+), 38 deletions(-)
 rename src/renderer/src/{providers/mcpToolUtils.ts => utils/mcp-tools.ts} (98%)

diff --git a/src/main/ipc.ts b/src/main/ipc.ts
index 0c80d560..ec163741 100644
--- a/src/main/ipc.ts
+++ b/src/main/ipc.ts
@@ -214,43 +214,23 @@ export function registerIpc(mainWindow: BrowserWindow, app: Electron.App) {
   )
 
   // Register MCP handlers
-  ipcMain.on('mcp:servers-from-renderer', (_event, servers) => {
-    mcpService.setServers(servers)
-  })
-
-  ipcMain.handle('mcp:list-servers', async () => {
-    return mcpService.listAvailableServices()
-  })
-
-  ipcMain.handle('mcp:add-server', async (_, server: MCPServer) => {
-    return mcpService.addServer(server)
-  })
-
-  ipcMain.handle('mcp:update-server', async (_, server: MCPServer) => {
-    return mcpService.updateServer(server)
-  })
-
-  ipcMain.handle('mcp:delete-server', async (_, serverName: string) => {
-    return mcpService.deleteServer(serverName)
-  })
-
-  ipcMain.handle('mcp:set-server-active', async (_, { name, isActive }) => {
-    return mcpService.setServerActive({ name, isActive })
-  })
+  ipcMain.on('mcp:servers-from-renderer', (_, servers) => mcpService.setServers(servers))
+  ipcMain.handle('mcp:list-servers', async () => mcpService.listAvailableServices())
+  ipcMain.handle('mcp:add-server', async (_, server: MCPServer) => mcpService.addServer(server))
+  ipcMain.handle('mcp:update-server', async (_, server: MCPServer) => mcpService.updateServer(server))
+  ipcMain.handle('mcp:delete-server', async (_, serverName: string) => mcpService.deleteServer(serverName))
+  ipcMain.handle('mcp:set-server-active', async (_, { name, isActive }) =>
+    mcpService.setServerActive({ name, isActive })
+  )
 
   // According to preload, this should take no parameters, but our implementation accepts
   // an optional serverName for better flexibility
-  ipcMain.handle('mcp:list-tools', async (_, serverName?: string) => {
-    return mcpService.listTools(serverName)
-  })
+  ipcMain.handle('mcp:list-tools', async (_, serverName?: string) => mcpService.listTools(serverName))
+  ipcMain.handle('mcp:call-tool', async (_, params: { client: string; name: string; args: any }) =>
+    mcpService.callTool(params)
+  )
 
-  ipcMain.handle('mcp:call-tool', async (_, params: { client: string; name: string; args: any }) => {
-    return mcpService.callTool(params)
-  })
-
-  ipcMain.handle('mcp:cleanup', async () => {
-    return mcpService.cleanup()
-  })
+  ipcMain.handle('mcp:cleanup', async () => mcpService.cleanup())
 
   // Listen for changes in MCP servers and notify renderer
   mcpService.on('servers-updated', (servers) => {
@@ -258,7 +238,5 @@ export function registerIpc(mainWindow: BrowserWindow, app: Electron.App) {
   })
 
   // Clean up MCP services when app quits
-  app.on('before-quit', async () => {
-    await mcpService.cleanup()
-  })
+  app.on('before-quit', () => mcpService.cleanup())
 }
diff --git a/src/renderer/src/pages/settings/MCPSettings.tsx b/src/renderer/src/pages/settings/MCPSettings.tsx
index 9975d577..7e9afe6c 100644
--- a/src/renderer/src/pages/settings/MCPSettings.tsx
+++ b/src/renderer/src/pages/settings/MCPSettings.tsx
@@ -134,6 +134,7 @@ const MCPSettings: FC = () => {
           window.api.mcp
             .addServer(mcpServer)
             .then(() => {
+              dispatch(addMCPServer(mcpServer))
               window.message.success(t('settings.mcp.addSuccess'))
               setLoading(false)
               setIsModalVisible(false)
@@ -143,7 +144,6 @@ const MCPSettings: FC = () => {
               window.message.error(`${t('settings.mcp.addError')}: ${error.message}`)
               setLoading(false)
             })
-          dispatch(addMCPServer(mcpServer))
         }
       })
       .catch(() => {
diff --git a/src/renderer/src/providers/mcpToolUtils.ts b/src/renderer/src/utils/mcp-tools.ts
similarity index 98%
rename from src/renderer/src/providers/mcpToolUtils.ts
rename to src/renderer/src/utils/mcp-tools.ts
index 3156a41e..d5e50652 100644
--- a/src/renderer/src/providers/mcpToolUtils.ts
+++ b/src/renderer/src/utils/mcp-tools.ts
@@ -3,7 +3,7 @@ import { FunctionCall, FunctionDeclaration, SchemaType, Tool as geminiToool } fr
 import { MCPServer, MCPTool, MCPToolResponse } from '@renderer/types'
 import { ChatCompletionMessageToolCall, ChatCompletionTool } from 'openai/resources'
 
-import { ChunkCallbackData } from '.'
+import { ChunkCallbackData } from '../providers'
 
 const supportedAttributes = [
   'type',