From b052e3280532ca3e10ec69e05417d8be1d0c71a8 Mon Sep 17 00:00:00 2001 From: Siddhant Rai Date: Wed, 10 Sep 2025 14:15:51 +0530 Subject: [PATCH] feat: enhance MCP tool configuration handling and authentication logic --- application/api/user/routes.py | 56 +++++++++++++++++- frontend/src/settings/ToolConfig.tsx | 88 ++++++++++++++++++++++++---- 2 files changed, 133 insertions(+), 11 deletions(-) diff --git a/application/api/user/routes.py b/application/api/user/routes.py index 2af52521..b0554461 100644 --- a/application/api/user/routes.py +++ b/application/api/user/routes.py @@ -3513,7 +3513,60 @@ class UpdateTool(Resource): ), 400, ) - update_data["config"] = data["config"] + 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"] user_tools_collection.update_one( @@ -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, diff --git a/frontend/src/settings/ToolConfig.tsx b/frontend/src/settings/ToolConfig.tsx index 61a1d850..bca5c6ce 100644 --- a/frontend/src/settings/ToolConfig.tsx +++ b/frontend/src/settings/ToolConfig.tsx @@ -30,9 +30,22 @@ export default function ToolConfig({ handleGoBack: () => void; }) { const token = useSelector(selectToken); - const [authKey, setAuthKey] = React.useState( - 'token' in tool.config ? tool.config.token : '', - ); + const [authKey, setAuthKey] = React.useState(() => { + 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( 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({
{Object.keys(tool?.config).length !== 0 && tool.name !== 'api_tool' && (

- {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')}

)}
@@ -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') + } />
)} @@ -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, },