From 10848f7a4580be0af150725114aa9807a2f5726b Mon Sep 17 00:00:00 2001
From: LiuVaayne <10231735+vaayne@users.noreply.github.com>
Date: Fri, 4 Apr 2025 09:57:54 +0800
Subject: [PATCH] feat: mcp tool permissions (#4348)

* feat(MCPSettings): add tool toggle functionality and update server configuration

* fix(McpSettings): improve server type handling and tool fetching logic
---
 .../settings/MCPSettings/McpSettings.tsx      | 66 +++++++++++++------
 .../pages/settings/MCPSettings/McpTool.tsx    | 47 +++++++++----
 src/renderer/src/services/ApiService.ts       |  4 +-
 src/renderer/src/types/index.ts               |  1 +
 4 files changed, 83 insertions(+), 35 deletions(-)

diff --git a/src/renderer/src/pages/settings/MCPSettings/McpSettings.tsx b/src/renderer/src/pages/settings/MCPSettings/McpSettings.tsx
index eb8f7fde..dfec53a1 100644
--- a/src/renderer/src/pages/settings/MCPSettings/McpSettings.tsx
+++ b/src/renderer/src/pages/settings/MCPSettings/McpSettings.tsx
@@ -93,22 +93,17 @@ const McpSettings: React.FC<Props> = ({ server }) => {
             .join('\n')
         : ''
     })
-    // eslint-disable-next-line react-hooks/exhaustive-deps
-  }, [server])
+  }, [server, form])
 
-  // Watch the serverType field to update the form layout dynamically
   useEffect(() => {
-    const type = form.getFieldValue('serverType')
-    type && setServerType(type)
-  }, [form])
-
-  // Load tools on initial mount if server is active
-  useEffect(() => {
-    fetchTools()
+    const currentServerType = form.getFieldValue('serverType')
+    if (currentServerType) {
+      setServerType(currentServerType)
+    }
     // eslint-disable-next-line react-hooks/exhaustive-deps
-  }, [])
+  }, [form.getFieldValue('serverType')])
 
-  const fetchTools = async () => {
+  const fetchTools = useCallback(async () => {
     if (server.isActive) {
       try {
         setLoadingServer(server.id)
@@ -124,7 +119,15 @@ const McpSettings: React.FC<Props> = ({ server }) => {
         setLoadingServer(null)
       }
     }
-  }
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [server.id])
+
+  useEffect(() => {
+    console.log('Loading tools for server:', server.id, 'Active:', server.isActive)
+    if (server.isActive) {
+      fetchTools()
+    }
+  }, [server.id, server.isActive, fetchTools])
 
   // Save the form data
   const onSave = async () => {
@@ -241,10 +244,6 @@ const McpSettings: React.FC<Props> = ({ server }) => {
     [server, t]
   )
 
-  const onFormValuesChange = () => {
-    setIsFormChanged(true)
-  }
-
   const formatError = (error: any) => {
     if (error.message.includes('32000')) {
       return t('settings.mcp.errors.32000')
@@ -278,6 +277,35 @@ const McpSettings: React.FC<Props> = ({ server }) => {
     }
   }
 
+  // Handle toggling a tool on/off
+  const handleToggleTool = useCallback(
+    async (tool: MCPTool, enabled: boolean) => {
+      // Create a new disabledTools array or use the existing one
+      let disabledTools = [...(server.disabledTools || [])]
+
+      if (enabled) {
+        // Remove tool from disabledTools if it's being enabled
+        disabledTools = disabledTools.filter((name) => name !== tool.name)
+      } else {
+        // Add tool to disabledTools if it's being disabled
+        if (!disabledTools.includes(tool.name)) {
+          disabledTools.push(tool.name)
+        }
+      }
+
+      // Update the server with new disabledTools
+      const updatedServer = {
+        ...server,
+        disabledTools
+      }
+
+      // Save the updated server configuration
+      // await window.api.mcp.updateServer(updatedServer)
+      updateMCPServer(updatedServer)
+    },
+    [server, updateMCPServer]
+  )
+
   return (
     <SettingContainer>
       <SettingGroup style={{ marginBottom: 0 }}>
@@ -302,7 +330,7 @@ const McpSettings: React.FC<Props> = ({ server }) => {
         <Form
           form={form}
           layout="vertical"
-          onValuesChange={onFormValuesChange}
+          onValuesChange={() => setIsFormChanged(true)}
           style={{
             // height: 'calc(100vh - var(--navbar-height) - 315px)',
             overflowY: 'auto',
@@ -380,7 +408,7 @@ const McpSettings: React.FC<Props> = ({ server }) => {
             </>
           )}
         </Form>
-        {server.isActive && <MCPToolsSection tools={tools} />}
+        {server.isActive && <MCPToolsSection tools={tools} server={server} onToggleTool={handleToggleTool} />}
       </SettingGroup>
     </SettingContainer>
   )
diff --git a/src/renderer/src/pages/settings/MCPSettings/McpTool.tsx b/src/renderer/src/pages/settings/MCPSettings/McpTool.tsx
index b5c80e89..2691c074 100644
--- a/src/renderer/src/pages/settings/MCPSettings/McpTool.tsx
+++ b/src/renderer/src/pages/settings/MCPSettings/McpTool.tsx
@@ -1,11 +1,27 @@
-import { MCPTool } from '@renderer/types'
-import { Badge, Collapse, Descriptions, Empty, Flex, Tag, Tooltip, Typography } from 'antd'
+import { MCPServer, MCPTool } from '@renderer/types'
+import { Badge, Collapse, Descriptions, Empty, Flex, Switch, Tag, Tooltip, Typography } from 'antd'
 import { useTranslation } from 'react-i18next'
 import styled from 'styled-components'
 
-const MCPToolsSection = ({ tools }: { tools: MCPTool[] }) => {
+interface MCPToolsSectionProps {
+  tools: MCPTool[]
+  server: MCPServer
+  onToggleTool: (tool: MCPTool, enabled: boolean) => void
+}
+
+const MCPToolsSection = ({ tools, server, onToggleTool }: MCPToolsSectionProps) => {
   const { t } = useTranslation()
 
+  // Check if a tool is enabled (not in the disabledTools array)
+  const isToolEnabled = (tool: MCPTool) => {
+    return !server.disabledTools?.includes(tool.name)
+  }
+
+  // Handle tool toggle
+  const handleToggle = (tool: MCPTool, checked: boolean) => {
+    onToggleTool(tool, checked)
+  }
+
   // Render tool properties from the input schema
   const renderToolProperties = (tool: MCPTool) => {
     if (!tool.inputSchema?.properties) return null
@@ -86,18 +102,21 @@ const MCPToolsSection = ({ tools }: { tools: MCPTool[] }) => {
             <Collapse.Panel
               key={tool.id}
               header={
-                <Flex vertical align="flex-start" style={{ width: '100%' }}>
-                  <Flex align="center" style={{ width: '100%' }}>
-                    <Typography.Text strong>{tool.name}</Typography.Text>
-                    <Typography.Text type="secondary" style={{ marginLeft: 8, fontSize: '12px' }}>
-                      {tool.id}
-                    </Typography.Text>
+                <Flex justify="space-between" align="center" style={{ width: '100%' }}>
+                  <Flex vertical align="flex-start">
+                    <Flex align="center" style={{ width: '100%' }}>
+                      <Typography.Text strong>{tool.name}</Typography.Text>
+                      <Typography.Text type="secondary" style={{ marginLeft: 8, fontSize: '12px' }}>
+                        {tool.id}
+                      </Typography.Text>
+                    </Flex>
+                    {tool.description && (
+                      <Typography.Text type="secondary" style={{ fontSize: '13px', marginTop: 4 }}>
+                        {tool.description}
+                      </Typography.Text>
+                    )}
                   </Flex>
-                  {tool.description && (
-                    <Typography.Text type="secondary" style={{ fontSize: '13px', marginTop: 4 }}>
-                      {tool.description}
-                    </Typography.Text>
-                  )}
+                  <Switch checked={isToolEnabled(tool)} onChange={(checked) => handleToggle(tool, checked)} />
                 </Flex>
               }>
               <SelectableContent>{renderToolProperties(tool)}</SelectableContent>
diff --git a/src/renderer/src/services/ApiService.ts b/src/renderer/src/services/ApiService.ts
index e70e5cba..459d9e82 100644
--- a/src/renderer/src/services/ApiService.ts
+++ b/src/renderer/src/services/ApiService.ts
@@ -106,8 +106,8 @@ export async function fetchChatCompletion({
     if (enabledMCPs && enabledMCPs.length > 0) {
       for (const mcpServer of enabledMCPs) {
         const tools = await window.api.mcp.listTools(mcpServer)
-        console.debug('tools', tools)
-        mcpTools.push(...tools)
+        const availableTools = tools.filter((tool: any) => !mcpServer.disabledTools?.includes(tool.name))
+        mcpTools.push(...availableTools)
       }
     }
 
diff --git a/src/renderer/src/types/index.ts b/src/renderer/src/types/index.ts
index cbac4a3e..0980bbe6 100644
--- a/src/renderer/src/types/index.ts
+++ b/src/renderer/src/types/index.ts
@@ -372,6 +372,7 @@ export interface MCPServer {
   args?: string[]
   env?: Record<string, string>
   isActive: boolean
+  disabledTools?: string[] // List of tool names that are disabled for this server
 }
 
 export interface MCPToolInputSchema {