🔧 fix: MCP Selection Persist and UI Flicker Issues (#9324)

* refactor: useMCPSelect

    - Add useGetMCPTools to use in useMCPSelect and elsewhere hooks for fetching MCP tools
    - remove memoized key
    - remove use of `useChatContext` and require conversationId as prop

* feat: Add MCPPanelContext and integrate conversationId as prop for useMCPSelect across components

- Introduced MCPPanelContext to manage conversationId state.
- Updated MCPSelect, MCPSubMenu, and MCPConfigDialog to accept conversationId as a prop.
- Modified ToolsDropdown and BadgeRow to pass conversationId to relevant components.
- Refactored MCPPanel to utilize MCPPanelProvider for context management.

* fix: remove nested ternary in ServerInitializationSection

- Replaced conditional operator with if-else statements for better readability in determining button text based on server initialization state and reinitialization status.

* refactor: wrap setValueWrap in useCallback for performance optimization

* refactor: streamline useMCPSelect by consolidating storageKey definition

* fix: prevent clearing selections on page refresh by tracking initial load completion

* refactor: simplify concern of useMCPSelect hook

* refactor: move ConfigFieldDetail interface to common types for better reusability, isolate usage of `useGetMCPTools`

* refactor: integrate mcpServerNames into BadgeRowContext and update ToolsDropdown and MCPSelect components
This commit is contained in:
Danny Avila 2025-08-28 00:44:49 -04:00 committed by GitHub
parent 2483623c88
commit c0511b9a5f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
19 changed files with 270 additions and 206 deletions

View file

@ -1 +1,3 @@
export * from './useMCPSelect';
export * from './useGetMCPTools';
export { useMCPServerManager } from './useMCPServerManager';

View file

@ -0,0 +1,43 @@
import { useMemo } from 'react';
import { Constants, EModelEndpoint } from 'librechat-data-provider';
import type { TPlugin } from 'librechat-data-provider';
import { useAvailableToolsQuery, useGetStartupConfig } from '~/data-provider';
export function useGetMCPTools() {
const { data: startupConfig } = useGetStartupConfig();
const { data: rawMcpTools } = useAvailableToolsQuery(EModelEndpoint.agents, {
select: (data: TPlugin[]) => {
const mcpToolsMap = new Map<string, TPlugin>();
data.forEach((tool) => {
const isMCP = tool.pluginKey.includes(Constants.mcp_delimiter);
if (isMCP) {
const parts = tool.pluginKey.split(Constants.mcp_delimiter);
const serverName = parts[parts.length - 1];
if (!mcpToolsMap.has(serverName)) {
mcpToolsMap.set(serverName, {
name: serverName,
pluginKey: tool.pluginKey,
authConfig: tool.authConfig,
authenticated: tool.authenticated,
});
}
}
});
return Array.from(mcpToolsMap.values());
},
});
const mcpToolDetails = useMemo(() => {
if (!rawMcpTools || !startupConfig?.mcpServers) {
return rawMcpTools;
}
return rawMcpTools.filter((tool) => {
const serverConfig = startupConfig?.mcpServers?.[tool.name];
return serverConfig?.chatMenu !== false;
});
}, [rawMcpTools, startupConfig?.mcpServers]);
return {
mcpToolDetails,
};
}

View file

@ -0,0 +1,72 @@
import { useRef, useCallback, useMemo } from 'react';
import { useRecoilState } from 'recoil';
import { Constants, LocalStorageKeys } from 'librechat-data-provider';
import useLocalStorage from '~/hooks/useLocalStorageAlt';
import { ephemeralAgentByConvoId } from '~/store';
const storageCondition = (value: unknown, rawCurrentValue?: string | null) => {
if (rawCurrentValue) {
try {
const currentValue = rawCurrentValue?.trim() ?? '';
if (currentValue.length > 2) {
return true;
}
} catch (e) {
console.error(e);
}
}
return Array.isArray(value) && value.length > 0;
};
export function useMCPSelect({ conversationId }: { conversationId?: string | null }) {
const key = conversationId ?? Constants.NEW_CONVO;
const [ephemeralAgent, setEphemeralAgent] = useRecoilState(ephemeralAgentByConvoId(key));
const storageKey = `${LocalStorageKeys.LAST_MCP_}${key}`;
const mcpState = useMemo(() => {
return ephemeralAgent?.mcp ?? [];
}, [ephemeralAgent?.mcp]);
const setSelectedValues = useCallback(
(values: string[] | null | undefined) => {
if (!values) {
return;
}
if (!Array.isArray(values)) {
return;
}
setEphemeralAgent((prev) => ({
...prev,
mcp: values,
}));
},
[setEphemeralAgent],
);
const [mcpValues, setMCPValuesRaw] = useLocalStorage<string[]>(
storageKey,
mcpState,
setSelectedValues,
storageCondition,
);
const setMCPValuesRawRef = useRef(setMCPValuesRaw);
setMCPValuesRawRef.current = setMCPValuesRaw;
/** Create a stable memoized setter to avoid re-creating it on every render and causing an infinite render loop */
const setMCPValues = useCallback((value: string[]) => {
setMCPValuesRawRef.current(value);
}, []);
const [isPinned, setIsPinned] = useLocalStorage<boolean>(
`${LocalStorageKeys.PIN_MCP_}${key}`,
true,
);
return {
isPinned,
mcpValues,
setIsPinned,
setMCPValues,
};
}

View file

@ -8,10 +8,10 @@ import {
useReinitializeMCPServerMutation,
} from 'librechat-data-provider/react-query';
import type { TUpdateUserPlugins, TPlugin } from 'librechat-data-provider';
import type { ConfigFieldDetail } from '~/components/MCP/MCPConfigDialog';
import type { ConfigFieldDetail } from '~/common';
import { useMCPConnectionStatusQuery } from '~/data-provider/Tools/queries';
import { useLocalize, useMCPSelect, useGetMCPTools } from '~/hooks';
import { useGetStartupConfig } from '~/data-provider';
import { useLocalize, useMCPSelect } from '~/hooks';
interface ServerState {
isInitializing: boolean;
@ -21,13 +21,14 @@ interface ServerState {
pollInterval: NodeJS.Timeout | null;
}
export function useMCPServerManager() {
export function useMCPServerManager({ conversationId }: { conversationId?: string | null }) {
const localize = useLocalize();
const { showToast } = useToastContext();
const mcpSelect = useMCPSelect();
const { data: startupConfig } = useGetStartupConfig();
const { mcpValues, setMCPValues, mcpToolDetails, isPinned, setIsPinned } = mcpSelect;
const queryClient = useQueryClient();
const { showToast } = useToastContext();
const { mcpToolDetails } = useGetMCPTools();
const mcpSelect = useMCPSelect({ conversationId });
const { data: startupConfig } = useGetStartupConfig();
const { mcpValues, setMCPValues, isPinned, setIsPinned } = mcpSelect;
const [isConfigModalOpen, setIsConfigModalOpen] = useState(false);
const [selectedToolForConfig, setSelectedToolForConfig] = useState<TPlugin | null>(null);
@ -90,7 +91,21 @@ export function useMCPServerManager() {
[connectionStatusData?.connectionStatus],
);
/** Filter disconnected servers when values change, but only after initial load
This prevents clearing selections on page refresh when servers haven't connected yet
*/
const hasInitialLoadCompleted = useRef(false);
useEffect(() => {
if (!connectionStatusData || Object.keys(connectionStatus).length === 0) {
return;
}
if (!hasInitialLoadCompleted.current) {
hasInitialLoadCompleted.current = true;
return;
}
if (!mcpValues?.length) return;
const connectedSelected = mcpValues.filter(
@ -100,7 +115,7 @@ export function useMCPServerManager() {
if (connectedSelected.length !== mcpValues.length) {
setMCPValues(connectedSelected);
}
}, [connectionStatus, mcpValues, setMCPValues]);
}, [connectionStatus, connectionStatusData, mcpValues, setMCPValues]);
const updateServerState = useCallback((serverName: string, updates: Partial<ServerState>) => {
setServerStates((prev) => {
@ -486,12 +501,12 @@ export function useMCPServerManager() {
};
},
[
isCancellable,
mcpToolDetails,
isInitializing,
cancelOAuthFlow,
connectionStatus,
startupConfig?.mcpServers,
isInitializing,
isCancellable,
cancelOAuthFlow,
],
);
@ -547,7 +562,6 @@ export function useMCPServerManager() {
mcpValues,
setMCPValues,
mcpToolDetails,
isPinned,
setIsPinned,
placeholderText,

View file

@ -1,4 +1,3 @@
export * from './useMCPSelect';
export * from './useToolToggle';
export { default as useAuthCodeTool } from './useAuthCodeTool';
export { default as usePluginInstall } from './usePluginInstall';

View file

@ -1,136 +0,0 @@
import { useRef, useEffect, useCallback, useMemo } from 'react';
import { useRecoilState } from 'recoil';
import { Constants, LocalStorageKeys, EModelEndpoint } from 'librechat-data-provider';
import type { TPlugin } from 'librechat-data-provider';
import { useAvailableToolsQuery, useGetStartupConfig } from '~/data-provider';
import useLocalStorage from '~/hooks/useLocalStorageAlt';
import { ephemeralAgentByConvoId } from '~/store';
import { useChatContext } from '~/Providers';
const storageCondition = (value: unknown, rawCurrentValue?: string | null) => {
if (rawCurrentValue) {
try {
const currentValue = rawCurrentValue?.trim() ?? '';
if (currentValue.length > 2) {
return true;
}
} catch (e) {
console.error(e);
}
}
return Array.isArray(value) && value.length > 0;
};
export function useMCPSelect() {
const { conversation } = useChatContext();
const key = useMemo(
() => conversation?.conversationId ?? Constants.NEW_CONVO,
[conversation?.conversationId],
);
const hasSetFetched = useRef<string | null>(null);
const [ephemeralAgent, setEphemeralAgent] = useRecoilState(ephemeralAgentByConvoId(key));
const { data: startupConfig } = useGetStartupConfig();
const { data: rawMcpTools, isFetched } = useAvailableToolsQuery(EModelEndpoint.agents, {
select: (data: TPlugin[]) => {
const mcpToolsMap = new Map<string, TPlugin>();
data.forEach((tool) => {
const isMCP = tool.pluginKey.includes(Constants.mcp_delimiter);
if (isMCP) {
const parts = tool.pluginKey.split(Constants.mcp_delimiter);
const serverName = parts[parts.length - 1];
if (!mcpToolsMap.has(serverName)) {
mcpToolsMap.set(serverName, {
name: serverName,
pluginKey: tool.pluginKey,
authConfig: tool.authConfig,
authenticated: tool.authenticated,
});
}
}
});
return Array.from(mcpToolsMap.values());
},
});
const mcpToolDetails = useMemo(() => {
if (!rawMcpTools || !startupConfig?.mcpServers) {
return rawMcpTools;
}
return rawMcpTools.filter((tool) => {
const serverConfig = startupConfig?.mcpServers?.[tool.name];
return serverConfig?.chatMenu !== false;
});
}, [rawMcpTools, startupConfig?.mcpServers]);
const mcpState = useMemo(() => {
return ephemeralAgent?.mcp ?? [];
}, [ephemeralAgent?.mcp]);
const setSelectedValues = useCallback(
(values: string[] | null | undefined) => {
if (!values) {
return;
}
if (!Array.isArray(values)) {
return;
}
setEphemeralAgent((prev) => ({
...prev,
mcp: values,
}));
},
[setEphemeralAgent],
);
const [mcpValues, setMCPValuesRaw] = useLocalStorage<string[]>(
`${LocalStorageKeys.LAST_MCP_}${key}`,
mcpState,
setSelectedValues,
storageCondition,
);
const setMCPValuesRawRef = useRef(setMCPValuesRaw);
setMCPValuesRawRef.current = setMCPValuesRaw;
// Create a stable memoized setter to avoid re-creating it on every render and causing an infinite render loop
const setMCPValues = useCallback((value: string[]) => {
setMCPValuesRawRef.current(value);
}, []);
const [isPinned, setIsPinned] = useLocalStorage<boolean>(
`${LocalStorageKeys.PIN_MCP_}${key}`,
true,
);
useEffect(() => {
if (hasSetFetched.current === key) {
return;
}
if (!isFetched) {
return;
}
hasSetFetched.current = key;
if ((mcpToolDetails?.length ?? 0) > 0) {
setMCPValues(mcpValues.filter((mcp) => mcpToolDetails?.some((tool) => tool.name === mcp)));
return;
}
setMCPValues([]);
}, [isFetched, setMCPValues, mcpToolDetails, key, mcpValues]);
const mcpServerNames = useMemo(() => {
return (mcpToolDetails ?? []).map((tool) => tool.name);
}, [mcpToolDetails]);
return {
isPinned,
mcpValues,
setIsPinned,
setMCPValues,
mcpServerNames,
ephemeralAgent,
mcpToolDetails,
setEphemeralAgent,
};
}

View file

@ -5,7 +5,7 @@
* - Also value will be updated everywhere, when value updated (via `storage` event)
*/
import { useEffect, useState } from 'react';
import { useEffect, useState, useCallback } from 'react';
export default function useLocalStorage<T>(
key: string,
@ -47,23 +47,26 @@ export default function useLocalStorage<T>(
// eslint-disable-next-line react-hooks/exhaustive-deps
}, [key, globalSetState]);
const setValueWrap = (value: T) => {
try {
setValue(value);
const storeLocal = () => {
localStorage.setItem(key, JSON.stringify(value));
window?.dispatchEvent(new StorageEvent('storage', { key }));
};
if (!storageCondition) {
storeLocal();
} else if (storageCondition(value, localStorage.getItem(key))) {
storeLocal();
const setValueWrap = useCallback(
(value: T) => {
try {
setValue(value);
const storeLocal = () => {
localStorage.setItem(key, JSON.stringify(value));
window?.dispatchEvent(new StorageEvent('storage', { key }));
};
if (!storageCondition) {
storeLocal();
} else if (storageCondition(value, localStorage.getItem(key))) {
storeLocal();
}
globalSetState?.(value);
} catch (e) {
console.error(e);
}
globalSetState?.(value);
} catch (e) {
console.error(e);
}
};
},
[key, globalSetState, storageCondition],
);
return [value, setValueWrap];
}