feat: enhance MCP tool configuration handling and authentication logic

application/api/user/routes.py

@@ -3513,6 +3513,59 @@ class UpdateTool(Resource):
                                 ),
                                 400,
                             )
+                tool_doc = user_tools_collection.find_one(
+                    {"_id": ObjectId(data["id"]), "user": user}
+                )
+                if tool_doc and tool_doc.get("name") == "mcp_tool":
+                    config = data["config"]
+                    existing_config = tool_doc.get("config", {})
+                    storage_config = existing_config.copy()
+
+                    storage_config.update(config)
+                    existing_credentials = {}
+                    if "encrypted_credentials" in existing_config:
+                        existing_credentials = decrypt_credentials(
+                            existing_config["encrypted_credentials"], user
+                        )
+                    auth_credentials = existing_credentials.copy()
+                    auth_type = storage_config.get("auth_type", "none")
+                    if auth_type == "api_key":
+                        if "api_key" in config and config["api_key"]:
+                            auth_credentials["api_key"] = config["api_key"]
+                        if "api_key_header" in config:
+                            auth_credentials["api_key_header"] = config[
+                                "api_key_header"
+                            ]
+                    elif auth_type == "bearer":
+                        if "bearer_token" in config and config["bearer_token"]:
+                            auth_credentials["bearer_token"] = config["bearer_token"]
+                        elif "encrypted_token" in config and config["encrypted_token"]:
+                            auth_credentials["bearer_token"] = config["encrypted_token"]
+                    elif auth_type == "basic":
+                        if "username" in config and config["username"]:
+                            auth_credentials["username"] = config["username"]
+                        if "password" in config and config["password"]:
+                            auth_credentials["password"] = config["password"]
+                    if auth_type != "none" and auth_credentials:
+                        encrypted_credentials_string = encrypt_credentials(
+                            auth_credentials, user
+                        )
+                        storage_config["encrypted_credentials"] = (
+                            encrypted_credentials_string
+                        )
+                    elif auth_type == "none":
+                        storage_config.pop("encrypted_credentials", None)
+                    for field in [
+                        "api_key",
+                        "bearer_token",
+                        "encrypted_token",
+                        "username",
+                        "password",
+                        "api_key_header",
+                    ]:
+                        storage_config.pop(field, None)
+                    update_data["config"] = storage_config
+                else:
                     update_data["config"] = data["config"]
             if "status" in data:
                 update_data["status"] = data["status"]
@@ -4238,6 +4291,7 @@ class MCPServerSave(Resource):
             tool_data = {
                 "name": "mcp_tool",
                 "displayName": data["displayName"],
+                "customName": data["displayName"],
                 "description": f"MCP Server: {storage_config.get('server_url', 'Unknown')}",
                 "config": storage_config,
                 "actions": transformed_actions,

frontend/src/settings/ToolConfig.tsx

@@ -30,9 +30,22 @@ export default function ToolConfig({
   handleGoBack: () => void;
 }) {
   const token = useSelector(selectToken);
-  const [authKey, setAuthKey] = React.useState<string>(
-    'token' in tool.config ? tool.config.token : '',
-  );
+  const [authKey, setAuthKey] = React.useState<string>(() => {
+    if (tool.name === 'mcp_tool') {
+      const config = tool.config as any;
+      if (config.auth_type === 'api_key') {
+        return config.api_key || '';
+      } else if (config.auth_type === 'bearer') {
+        return config.encrypted_token || '';
+      } else if (config.auth_type === 'basic') {
+        return config.password || '';
+      }
+      return '';
+    } else if ('token' in tool.config) {
+      return tool.config.token;
+    }
+    return '';
+  });
   const [customName, setCustomName] = React.useState<string>(
     tool.customName || '',
   );
@@ -97,6 +110,26 @@ export default function ToolConfig({
   };
 
   const handleSaveChanges = () => {
+    let configToSave;
+    if (tool.name === 'api_tool') {
+      configToSave = tool.config;
+    } else if (tool.name === 'mcp_tool') {
+      configToSave = { ...tool.config } as any;
+      const mcpConfig = tool.config as any;
+
+      if (authKey.trim()) {
+        if (mcpConfig.auth_type === 'api_key') {
+          configToSave.api_key = authKey;
+        } else if (mcpConfig.auth_type === 'bearer') {
+          configToSave.encrypted_token = authKey;
+        } else if (mcpConfig.auth_type === 'basic') {
+          configToSave.password = authKey;
+        }
+      }
+    } else {
+      configToSave = { token: authKey };
+    }
+
     userService
       .updateTool(
         {
@@ -105,7 +138,7 @@ export default function ToolConfig({
           displayName: tool.displayName,
           customName: customName,
           description: tool.description,
-          config: tool.name === 'api_tool' ? tool.config : { token: authKey },
+          config: configToSave,
           actions: 'actions' in tool ? tool.actions : [],
           status: tool.status,
         },
@@ -196,7 +229,15 @@ export default function ToolConfig({
       <div className="mt-1">
         {Object.keys(tool?.config).length !== 0 && tool.name !== 'api_tool' && (
           <p className="text-eerie-black dark:text-bright-gray text-sm font-semibold">
-            {t('settings.tools.authentication')}
+            {tool.name === 'mcp_tool'
+              ? (tool.config as any)?.auth_type === 'bearer'
+                ? 'Bearer Token'
+                : (tool.config as any)?.auth_type === 'api_key'
+                  ? 'API Key'
+                  : (tool.config as any)?.auth_type === 'basic'
+                    ? 'Password'
+                    : t('settings.tools.authentication')
+              : t('settings.tools.authentication')}
           </p>
         )}
         <div className="mt-4 flex flex-col items-start gap-2 sm:flex-row sm:items-center">
@@ -208,7 +249,17 @@ export default function ToolConfig({
                   value={authKey}
                   onChange={(e) => setAuthKey(e.target.value)}
                   borderVariant="thin"
-                  placeholder={t('modals.configTool.apiKeyPlaceholder')}
+                  placeholder={
+                    tool.name === 'mcp_tool'
+                      ? (tool.config as any)?.auth_type === 'bearer'
+                        ? 'Bearer Token'
+                        : (tool.config as any)?.auth_type === 'api_key'
+                          ? 'API Key'
+                          : (tool.config as any)?.auth_type === 'basic'
+                            ? 'Password'
+                            : t('modals.configTool.apiKeyPlaceholder')
+                      : t('modals.configTool.apiKeyPlaceholder')
+                  }
                 />
               </div>
             )}
@@ -450,6 +501,26 @@ export default function ToolConfig({
             setModalState={(state) => setShowUnsavedModal(state === 'ACTIVE')}
             submitLabel={t('settings.tools.saveAndLeave')}
             handleSubmit={() => {
+              let configToSave;
+              if (tool.name === 'api_tool') {
+                configToSave = tool.config;
+              } else if (tool.name === 'mcp_tool') {
+                configToSave = { ...tool.config } as any;
+                const mcpConfig = tool.config as any;
+
+                if (authKey.trim()) {
+                  if (mcpConfig.auth_type === 'api_key') {
+                    configToSave.api_key = authKey;
+                  } else if (mcpConfig.auth_type === 'bearer') {
+                    configToSave.encrypted_token = authKey;
+                  } else if (mcpConfig.auth_type === 'basic') {
+                    configToSave.password = authKey;
+                  }
+                }
+              } else {
+                configToSave = { token: authKey };
+              }
+
               userService
                 .updateTool(
                   {
@@ -458,10 +529,7 @@ export default function ToolConfig({
                     displayName: tool.displayName,
                     customName: customName,
                     description: tool.description,
-                    config:
-                      tool.name === 'api_tool'
-                        ? tool.config
-                        : { token: authKey },
+                    config: configToSave,
                     actions: 'actions' in tool ? tool.actions : [],
                     status: tool.status,
                   },