From a6bf2b6ce342b87f38aabdf99f36e88f5e7b7f31 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Mon, 22 Sep 2025 08:53:19 -0400 Subject: [PATCH] =?UTF-8?q?=F0=9F=94=90=20refactor:=20Improve=20MCP=20Auth?= =?UTF-8?q?=20UX=20for=20Agent=20Panel=20(#9762)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * chore: Improve logging format for initial flow state creation * refactor: MCP Tool Management with Improved Auth Handling and State Management - Updated `CustomUserVarsSection` to include optional localization for placeholder text. - Refactored `MCPToolSelectDialog` to streamline tool addition and management, including improved handling of authentication data. - Introduced a new `addToolsToForm` function to encapsulate logic for adding tools to the form state. - Enhanced `useRemoveMCPTool` hook to simplify tool removal logic and ensure proper state updates. - Added loading state management for custom variable saving to improve user experience. * refactor: Enhance MCP Tool Removal Logic and Integrate Toast Notifications - Updated `MCPToolSelectDialog` to utilize the new `removeTool` function from the `useRemoveMCPTool` hook for improved tool removal handling. - Refactored `useRemoveMCPTool` to accept options for toast notifications, allowing for more flexible user feedback during tool removal. - Removed the previous inline tool removal logic to streamline the component's code and improve maintainability. * refactor: Enhance user plugins mutation to invalidate MCP auth values on uninstall * refactor: Replace refetchQueries with invalidateQueries for improved cache management * chore: remove unused i18n key --- .../components/MCP/CustomUserVarsSection.tsx | 2 +- .../src/components/SidePanel/MCP/MCPPanel.tsx | 6 +- .../components/Tools/MCPToolSelectDialog.tsx | 150 ++++++++++-------- client/src/hooks/MCP/useMCPServerManager.ts | 6 +- client/src/hooks/MCP/useRemoveMCPTool.ts | 52 +++--- client/src/locales/en/translation.json | 1 - packages/api/src/flow/manager.ts | 2 +- .../src/react-query/react-query-service.ts | 6 +- 8 files changed, 111 insertions(+), 114 deletions(-) diff --git a/client/src/components/MCP/CustomUserVarsSection.tsx b/client/src/components/MCP/CustomUserVarsSection.tsx index ded7115b15..9c2f0870d4 100644 --- a/client/src/components/MCP/CustomUserVarsSection.tsx +++ b/client/src/components/MCP/CustomUserVarsSection.tsx @@ -66,7 +66,7 @@ function AuthField({ name, config, hasValue, control, errors }: AuthFieldProps) placeholder={ hasValue ? localize('com_ui_mcp_update_var', { 0: config.title }) - : localize('com_ui_mcp_enter_var', { 0: config.title }) + : `${localize('com_ui_mcp_enter_var', { 0: config.title })} ${localize('com_ui_optional')}` } className="w-full rounded border border-border-medium bg-transparent px-2 py-1 text-text-primary placeholder:text-text-secondary focus:outline-none sm:text-sm" /> diff --git a/client/src/components/SidePanel/MCP/MCPPanel.tsx b/client/src/components/SidePanel/MCP/MCPPanel.tsx index 8f79023897..d189478f1d 100644 --- a/client/src/components/SidePanel/MCP/MCPPanel.tsx +++ b/client/src/components/SidePanel/MCP/MCPPanel.tsx @@ -31,9 +31,9 @@ function MCPPanelContent() { showToast({ message: localize('com_nav_mcp_vars_updated'), status: 'success' }); await Promise.all([ - queryClient.refetchQueries([QueryKeys.mcpTools]), - queryClient.refetchQueries([QueryKeys.mcpAuthValues]), - queryClient.refetchQueries([QueryKeys.mcpConnectionStatus]), + queryClient.invalidateQueries([QueryKeys.mcpTools]), + queryClient.invalidateQueries([QueryKeys.mcpAuthValues]), + queryClient.invalidateQueries([QueryKeys.mcpConnectionStatus]), ]); }, onError: (error: unknown) => { diff --git a/client/src/components/Tools/MCPToolSelectDialog.tsx b/client/src/components/Tools/MCPToolSelectDialog.tsx index 95588826d6..6e3b75f099 100644 --- a/client/src/components/Tools/MCPToolSelectDialog.tsx +++ b/client/src/components/Tools/MCPToolSelectDialog.tsx @@ -1,12 +1,18 @@ import { useEffect, useState, useMemo } from 'react'; import { Search, X } from 'lucide-react'; import { useFormContext } from 'react-hook-form'; -import { Constants, EModelEndpoint } from 'librechat-data-provider'; +import { useQueryClient } from '@tanstack/react-query'; +import { Constants, EModelEndpoint, QueryKeys } from 'librechat-data-provider'; import { Dialog, DialogPanel, DialogTitle, Description } from '@headlessui/react'; import { useUpdateUserPluginsMutation } from 'librechat-data-provider/react-query'; import type { TError, AgentToolType } from 'librechat-data-provider'; import type { AgentForm, TPluginStoreDialogProps } from '~/common'; -import { useLocalize, usePluginDialogHelpers, useMCPServerManager } from '~/hooks'; +import { + usePluginDialogHelpers, + useMCPServerManager, + useRemoveMCPTool, + useLocalize, +} from '~/hooks'; import CustomUserVarsSection from '~/components/MCP/CustomUserVarsSection'; import { PluginPagination } from '~/components/Plugins/Store'; import { useAgentPanelContext } from '~/Providers'; @@ -24,13 +30,16 @@ function MCPToolSelectDialog({ endpoint: EModelEndpoint.agents; }) { const localize = useLocalize(); + const queryClient = useQueryClient(); const { initializeServer } = useMCPServerManager(); const { getValues, setValue } = useFormContext(); + const { removeTool } = useRemoveMCPTool({ showToast: false }); const { mcpServersMap, startupConfig } = useAgentPanelContext(); const { refetch: refetchMCPTools } = useMCPToolsQuery({ enabled: mcpServersMap.size > 0, }); + const [isSavingCustomVars, setIsSavingCustomVars] = useState(false); const [isInitializing, setIsInitializing] = useState(null); const [configuringServer, setConfiguringServer] = useState(null); @@ -67,9 +76,27 @@ function MCPToolSelectDialog({ }, 5000); }; - const handleDirectAdd = async (serverName: string) => { + const handleDirectAdd = async (serverName: string, authData?: Record) => { try { setIsInitializing(serverName); + + // First, save auth if provided + if (authData && Object.keys(authData).length > 0) { + await updateUserPlugins.mutateAsync({ + pluginKey: `${Constants.mcp_prefix}${serverName}`, + action: 'install', + auth: authData, + isEntityTool: true, + }); + + // Invalidate auth values query to ensure fresh data + await queryClient.invalidateQueries([QueryKeys.mcpAuthValues, serverName]); + + // Small delay to ensure backend has processed the auth + await new Promise((resolve) => setTimeout(resolve, 500)); + } + + // Then initialize server if needed const serverInfo = mcpServersMap.get(serverName); if (!serverInfo?.isConnected) { const result = await initializeServer(serverName); @@ -78,64 +105,65 @@ function MCPToolSelectDialog({ return; } } - updateUserPlugins.mutate( - { - pluginKey: `${Constants.mcp_prefix}${serverName}`, - action: 'install', - auth: {}, - isEntityTool: true, - }, - { - onError: (error: unknown) => { - handleInstallError(error as TError); - setIsInitializing(null); - }, - onSuccess: async () => { - const { data: updatedMCPData } = await refetchMCPTools(); - const currentTools = getValues('tools') || []; - const toolsToAdd: string[] = [ - `${Constants.mcp_server}${Constants.mcp_delimiter}${serverName}`, - ]; - - if (updatedMCPData?.servers?.[serverName]) { - const serverData = updatedMCPData.servers[serverName]; - serverData.tools.forEach((tool) => { - toolsToAdd.push(tool.pluginKey); - }); - } - - const newTools = toolsToAdd.filter((tool) => !currentTools.includes(tool)); - if (newTools.length > 0) { - setValue('tools', [...currentTools, ...newTools]); - } - setIsInitializing(null); - }, - }, - ); + // Finally, add tools to form + await addToolsToForm(serverName); + setIsInitializing(null); } catch (error) { console.error('Error adding MCP server:', error); + handleInstallError(error as TError); + setIsInitializing(null); + } + }; + + const addToolsToForm = async (serverName: string) => { + const { data: updatedMCPData } = await refetchMCPTools(); + + const currentTools = getValues('tools') || []; + const toolsToAdd: string[] = [`${Constants.mcp_server}${Constants.mcp_delimiter}${serverName}`]; + + if (updatedMCPData?.servers?.[serverName]) { + const serverData = updatedMCPData.servers[serverName]; + serverData.tools.forEach((tool) => { + toolsToAdd.push(tool.pluginKey); + }); + } + + const newTools = toolsToAdd.filter((tool) => !currentTools.includes(tool)); + if (newTools.length > 0) { + setValue('tools', [...currentTools, ...newTools]); } }; const handleSaveCustomVars = async (serverName: string, authData: Record) => { try { - await updateUserPlugins.mutateAsync({ - pluginKey: `${Constants.mcp_prefix}${serverName}`, - action: 'install', - auth: authData, - isEntityTool: true, + setIsSavingCustomVars(true); + + // Filter out empty values to avoid overwriting existing values with empty ones + const filteredAuthData: Record = {}; + Object.entries(authData).forEach(([key, value]) => { + if (value && value.trim()) { + filteredAuthData[key] = value.trim(); + } }); - await handleDirectAdd(serverName); - + // Always add the tool, but only pass auth data if there are values to save + // Empty auth data is fine - the tool can work without credentials + await handleDirectAdd( + serverName, + Object.keys(filteredAuthData).length > 0 ? filteredAuthData : undefined, + ); setConfiguringServer(null); } catch (error) { console.error('Error saving custom vars:', error); + handleInstallError(error as TError); + } finally { + setIsSavingCustomVars(false); } }; const handleRevokeCustomVars = (serverName: string) => { + setIsSavingCustomVars(true); updateUserPlugins.mutate( { pluginKey: `${Constants.mcp_prefix}${serverName}`, @@ -144,9 +172,13 @@ function MCPToolSelectDialog({ isEntityTool: true, }, { - onError: (error: unknown) => handleInstallError(error as TError), - onSuccess: () => { + onError: (error: unknown) => { + handleInstallError(error as TError); + setIsSavingCustomVars(false); + }, + onSuccess: async () => { setConfiguringServer(null); + setIsSavingCustomVars(false); }, }, ); @@ -170,28 +202,6 @@ function MCPToolSelectDialog({ } }; - const onRemoveTool = (serverName: string) => { - updateUserPlugins.mutate( - { - pluginKey: `${Constants.mcp_prefix}${serverName}`, - action: 'uninstall', - auth: {}, - isEntityTool: true, - }, - { - onError: (error: unknown) => handleInstallError(error as TError), - onSuccess: () => { - const currentTools = getValues('tools') || []; - const remainingTools = currentTools.filter( - (tool) => - tool !== serverName && !tool.endsWith(`${Constants.mcp_delimiter}${serverName}`), - ); - setValue('tools', remainingTools); - }, - }, - ); - }; - const installedToolsSet = useMemo(() => { return new Set(mcpServerNames); }, [mcpServerNames]); @@ -289,10 +299,10 @@ function MCPToolSelectDialog({ handleSaveCustomVars(configuringServer, authData)} onRevoke={() => handleRevokeCustomVars(configuringServer)} - isSubmitting={updateUserPlugins.isLoading} /> )} @@ -342,7 +352,7 @@ function MCPToolSelectDialog({ isConfiguring={isConfiguring} isInitializing={isServerInitializing} onAddTool={() => onAddTool(serverInfo.serverName)} - onRemoveTool={() => onRemoveTool(serverInfo.serverName)} + onRemoveTool={() => removeTool(serverInfo.serverName)} /> ); })} diff --git a/client/src/hooks/MCP/useMCPServerManager.ts b/client/src/hooks/MCP/useMCPServerManager.ts index 2e05569710..101904d29b 100644 --- a/client/src/hooks/MCP/useMCPServerManager.ts +++ b/client/src/hooks/MCP/useMCPServerManager.ts @@ -52,9 +52,9 @@ export function useMCPServerManager({ conversationId }: { conversationId?: strin showToast({ message: localize('com_nav_mcp_vars_updated'), status: 'success' }); await Promise.all([ - queryClient.refetchQueries([QueryKeys.mcpTools]), - queryClient.refetchQueries([QueryKeys.mcpAuthValues]), - queryClient.refetchQueries([QueryKeys.mcpConnectionStatus]), + queryClient.invalidateQueries([QueryKeys.mcpTools]), + queryClient.invalidateQueries([QueryKeys.mcpAuthValues]), + queryClient.invalidateQueries([QueryKeys.mcpConnectionStatus]), ]); }, onError: (error: unknown) => { diff --git a/client/src/hooks/MCP/useRemoveMCPTool.ts b/client/src/hooks/MCP/useRemoveMCPTool.ts index af14f5c084..1f35e25c37 100644 --- a/client/src/hooks/MCP/useRemoveMCPTool.ts +++ b/client/src/hooks/MCP/useRemoveMCPTool.ts @@ -2,19 +2,19 @@ import { useCallback } from 'react'; import { useFormContext } from 'react-hook-form'; import { Constants } from 'librechat-data-provider'; import { useToastContext } from '@librechat/client'; -import { useUpdateUserPluginsMutation } from 'librechat-data-provider/react-query'; import type { AgentForm } from '~/common'; import { useLocalize } from '~/hooks'; /** * Hook for removing MCP tools/servers from an agent * Provides unified logic for MCPTool, UninitializedMCPTool, and UnconfiguredMCPTool components + * Note: This only removes the tool from the form, it does not delete associated auth credentials */ -export function useRemoveMCPTool() { +export function useRemoveMCPTool(options?: { showToast?: boolean }) { const localize = useLocalize(); const { showToast } = useToastContext(); - const updateUserPlugins = useUpdateUserPluginsMutation(); const { getValues, setValue } = useFormContext(); + const shouldShowToast = options?.showToast !== false; const removeTool = useCallback( (serverName: string) => { @@ -22,39 +22,23 @@ export function useRemoveMCPTool() { return; } - updateUserPlugins.mutate( - { - pluginKey: `${Constants.mcp_prefix}${serverName}`, - action: 'uninstall', - auth: {}, - isEntityTool: true, - }, - { - onError: (error: unknown) => { - showToast({ - message: localize('com_ui_delete_tool_error', { error: String(error) }), - status: 'error', - }); - }, - onSuccess: () => { - const currentTools = getValues('tools'); - const remainingToolIds = - currentTools?.filter( - (currentToolId) => - currentToolId !== serverName && - !currentToolId.endsWith(`${Constants.mcp_delimiter}${serverName}`), - ) || []; - setValue('tools', remainingToolIds, { shouldDirty: true }); + const currentTools = getValues('tools'); + const remainingToolIds = + currentTools?.filter( + (currentToolId) => + currentToolId !== serverName && + !currentToolId.endsWith(`${Constants.mcp_delimiter}${serverName}`), + ) || []; + setValue('tools', remainingToolIds, { shouldDirty: true }); - showToast({ - message: localize('com_ui_delete_tool_save_reminder'), - status: 'warning', - }); - }, - }, - ); + if (shouldShowToast) { + showToast({ + message: localize('com_ui_delete_tool_save_reminder'), + status: 'warning', + }); + } }, - [getValues, setValue, updateUserPlugins, showToast, localize], + [getValues, setValue, showToast, localize, shouldShowToast], ); return { removeTool }; diff --git a/client/src/locales/en/translation.json b/client/src/locales/en/translation.json index 7c2e9e7f3e..b1bc9ba96e 100644 --- a/client/src/locales/en/translation.json +++ b/client/src/locales/en/translation.json @@ -834,7 +834,6 @@ "com_ui_delete_success": "Successfully deleted", "com_ui_delete_tool": "Delete Tool", "com_ui_delete_tool_confirm": "Are you sure you want to delete this tool?", - "com_ui_delete_tool_error": "Error while deleting the tool: {{error}}", "com_ui_delete_tool_save_reminder": "Tool removed. Save the agent to apply changes.", "com_ui_deleted": "Deleted", "com_ui_deleting_file": "Deleting file...", diff --git a/packages/api/src/flow/manager.ts b/packages/api/src/flow/manager.ts index 09ef40ec81..af7a3e9112 100644 --- a/packages/api/src/flow/manager.ts +++ b/packages/api/src/flow/manager.ts @@ -74,7 +74,7 @@ export class FlowStateManager { createdAt: Date.now(), }; - logger.debug('Creating initial flow state:', flowKey); + logger.debug(`[${flowKey}] Creating initial flow state`); await this.keyv.set(flowKey, initialState, this.ttl); return this.monitorFlow(flowKey, type, signal); } diff --git a/packages/data-provider/src/react-query/react-query-service.ts b/packages/data-provider/src/react-query/react-query-service.ts index db1f0c9052..c4b55e2039 100644 --- a/packages/data-provider/src/react-query/react-query-service.ts +++ b/packages/data-provider/src/react-query/react-query-service.ts @@ -320,6 +320,10 @@ export const useUpdateUserPluginsMutation = ( onSuccess: (...args) => { queryClient.invalidateQueries([QueryKeys.user]); onSuccess?.(...args); + if (args[1]?.action === 'uninstall' && args[1]?.pluginKey?.startsWith(Constants.mcp_prefix)) { + const serverName = args[1]?.pluginKey?.substring(Constants.mcp_prefix.length); + queryClient.invalidateQueries([QueryKeys.mcpAuthValues, serverName]); + } }, }); }; @@ -339,7 +343,7 @@ export const useReinitializeMCPServerMutation = (): UseMutationResult< const queryClient = useQueryClient(); return useMutation((serverName: string) => dataService.reinitializeMCPServer(serverName), { onSuccess: () => { - queryClient.refetchQueries([QueryKeys.mcpTools]); + queryClient.invalidateQueries([QueryKeys.mcpTools]); }, }); };