mirror of
https://github.com/danny-avila/LibreChat.git
synced 2025-12-16 16:30:15 +01:00
🔐 refactor: Improve MCP Auth UX for Agent Panel (#9762)
* 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
This commit is contained in:
Danny Avila
GitHub
parent
ff8dac570f
commit
a6bf2b6ce3
8 changed files with 111 additions and 114 deletions
|
|
@ -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"
|
||||
/>
|
||||
|
|
|
|||
|
|
@ -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) => {
|
||||
|
|
|
|||
|
|
@ -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<AgentForm>();
|
||||
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<string | null>(null);
|
||||
const [configuringServer, setConfiguringServer] = useState<string | null>(null);
|
||||
|
||||
|
|
@ -67,9 +76,27 @@ function MCPToolSelectDialog({
|
|||
}, 5000);
|
||||
};
|
||||
|
||||
const handleDirectAdd = async (serverName: string) => {
|
||||
const handleDirectAdd = async (serverName: string, authData?: Record<string, string>) => {
|
||||
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<string, string>) => {
|
||||
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<string, string> = {};
|
||||
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({
|
|||
</div>
|
||||
<CustomUserVarsSection
|
||||
serverName={configuringServer}
|
||||
isSubmitting={isSavingCustomVars}
|
||||
fields={startupConfig?.mcpServers?.[configuringServer]?.customUserVars || {}}
|
||||
onSave={(authData) => handleSaveCustomVars(configuringServer, authData)}
|
||||
onRevoke={() => handleRevokeCustomVars(configuringServer)}
|
||||
isSubmitting={updateUserPlugins.isLoading}
|
||||
/>
|
||||
</div>
|
||||
)}
|
||||
|
|
@ -342,7 +352,7 @@ function MCPToolSelectDialog({
|
|||
isConfiguring={isConfiguring}
|
||||
isInitializing={isServerInitializing}
|
||||
onAddTool={() => onAddTool(serverInfo.serverName)}
|
||||
onRemoveTool={() => onRemoveTool(serverInfo.serverName)}
|
||||
onRemoveTool={() => removeTool(serverInfo.serverName)}
|
||||
/>
|
||||
);
|
||||
})}
|
||||
|
|
|
|||
|
|
@ -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) => {
|
||||
|
|
|
|||
|
|
@ -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<AgentForm>();
|
||||
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 };
|
||||
|
|
|
|||
|
|
@ -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...",
|
||||
|
|
|
|||
|
|
@ -74,7 +74,7 @@ export class FlowStateManager<T = unknown> {
|
|||
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);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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]);
|
||||
},
|
||||
});
|
||||
};
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue