mirror of
https://github.com/danny-avila/LibreChat.git
synced 2025-12-19 09:50:15 +01:00
🏹 feat: Concurrent MCP Initialization Support (#8677)
* ✨ feat: Enhance MCP Connection Status Management - Introduced new functions to retrieve and manage connection status for multiple MCP servers, including OAuth flow checks and server-specific status retrieval. - Refactored the MCP connection status endpoints to support both all servers and individual server queries. - Replaced the old server initialization hook with a new `useMCPServerManager` hook for improved state management and handling of multiple OAuth flows. - Updated the MCPPanel component to utilize the new context provider for better state handling and UI updates. - Fixed a number of UI bugs when initializing servers * 🗣️ i18n: Remove unused strings from translation.json * refactor: move helper functions out of the route module into mcp service file * ci: add tests for newly added functions in mcp service file * fix: memoize setMCPValues to avoid render loop
This commit is contained in:
Dustin Healy
GitHub
parent
37aba18a96
commit
0ef3fefaec
15 changed files with 1092 additions and 542 deletions
|
|
@ -1 +1 @@
|
|||
export { useMCPServerInitialization } from './useMCPServerInitialization';
|
||||
export { useMCPServerManager } from './useMCPServerManager';
|
||||
|
|
|
|||
|
|
@ -1,317 +0,0 @@
|
|||
import { useCallback, useState, useEffect, useMemo } from 'react';
|
||||
import { useToastContext } from '@librechat/client';
|
||||
import { QueryKeys } from 'librechat-data-provider';
|
||||
import { useQueryClient } from '@tanstack/react-query';
|
||||
import {
|
||||
useReinitializeMCPServerMutation,
|
||||
useCancelMCPOAuthMutation,
|
||||
} from 'librechat-data-provider/react-query';
|
||||
import { useMCPConnectionStatusQuery } from '~/data-provider/Tools/queries';
|
||||
import { useLocalize } from '~/hooks';
|
||||
import { logger } from '~/utils';
|
||||
|
||||
interface UseMCPServerInitializationOptions {
|
||||
onSuccess?: (serverName: string) => void;
|
||||
onOAuthStarted?: (serverName: string, oauthUrl: string) => void;
|
||||
onError?: (serverName: string, error: any) => void;
|
||||
}
|
||||
|
||||
export function useMCPServerInitialization(options?: UseMCPServerInitializationOptions) {
|
||||
const localize = useLocalize();
|
||||
const { showToast } = useToastContext();
|
||||
const queryClient = useQueryClient();
|
||||
|
||||
// OAuth state management
|
||||
const [oauthPollingServers, setOauthPollingServers] = useState<Map<string, string>>(new Map());
|
||||
const [oauthStartTimes, setOauthStartTimes] = useState<Map<string, number>>(new Map());
|
||||
const [initializingServers, setInitializingServers] = useState<Set<string>>(new Set());
|
||||
const [cancellableServers, setCancellableServers] = useState<Set<string>>(new Set());
|
||||
|
||||
// Get connection status
|
||||
const { data: connectionStatusData } = useMCPConnectionStatusQuery();
|
||||
const connectionStatus = useMemo(
|
||||
() => connectionStatusData?.connectionStatus || {},
|
||||
[connectionStatusData],
|
||||
);
|
||||
|
||||
// Main initialization mutation
|
||||
const reinitializeMutation = useReinitializeMCPServerMutation();
|
||||
|
||||
// Track which server is currently being processed
|
||||
const [currentProcessingServer, setCurrentProcessingServer] = useState<string | null>(null);
|
||||
|
||||
// Cancel OAuth mutation
|
||||
const cancelOAuthMutation = useCancelMCPOAuthMutation();
|
||||
|
||||
// Helper function to clean up OAuth state
|
||||
const cleanupOAuthState = useCallback((serverName: string) => {
|
||||
setOauthPollingServers((prev) => {
|
||||
const newMap = new Map(prev);
|
||||
newMap.delete(serverName);
|
||||
return newMap;
|
||||
});
|
||||
|
||||
setOauthStartTimes((prev) => {
|
||||
const newMap = new Map(prev);
|
||||
newMap.delete(serverName);
|
||||
return newMap;
|
||||
});
|
||||
|
||||
setInitializingServers((prev) => {
|
||||
const newSet = new Set(prev);
|
||||
newSet.delete(serverName);
|
||||
return newSet;
|
||||
});
|
||||
|
||||
setCancellableServers((prev) => {
|
||||
const newSet = new Set(prev);
|
||||
newSet.delete(serverName);
|
||||
return newSet;
|
||||
});
|
||||
}, []);
|
||||
|
||||
// Cancel OAuth flow
|
||||
const cancelOAuthFlow = useCallback(
|
||||
(serverName: string) => {
|
||||
logger.info(`[MCP OAuth] User cancelling OAuth flow for ${serverName}`);
|
||||
|
||||
cancelOAuthMutation.mutate(serverName, {
|
||||
onSuccess: () => {
|
||||
cleanupOAuthState(serverName);
|
||||
showToast({
|
||||
message: localize('com_ui_mcp_oauth_cancelled', { 0: serverName }),
|
||||
status: 'info',
|
||||
});
|
||||
},
|
||||
onError: (error) => {
|
||||
logger.error(`[MCP OAuth] Failed to cancel OAuth flow for ${serverName}:`, error);
|
||||
// Clean up state anyway
|
||||
cleanupOAuthState(serverName);
|
||||
},
|
||||
});
|
||||
},
|
||||
[cancelOAuthMutation, cleanupOAuthState, showToast, localize],
|
||||
);
|
||||
|
||||
// Helper function to handle successful connection
|
||||
const handleSuccessfulConnection = useCallback(
|
||||
async (serverName: string, message: string) => {
|
||||
showToast({ message, status: 'success' });
|
||||
|
||||
// Force immediate refetch to update UI
|
||||
await Promise.all([
|
||||
queryClient.refetchQueries([QueryKeys.mcpConnectionStatus]),
|
||||
queryClient.refetchQueries([QueryKeys.tools]),
|
||||
]);
|
||||
|
||||
// Clean up OAuth state
|
||||
cleanupOAuthState(serverName);
|
||||
|
||||
// Call optional success callback
|
||||
options?.onSuccess?.(serverName);
|
||||
},
|
||||
[showToast, queryClient, options, cleanupOAuthState],
|
||||
);
|
||||
|
||||
// Helper function to handle OAuth timeout/failure
|
||||
const handleOAuthFailure = useCallback(
|
||||
(serverName: string, isTimeout: boolean) => {
|
||||
logger.warn(
|
||||
`[MCP OAuth] OAuth ${isTimeout ? 'timed out' : 'failed'} for ${serverName}, stopping poll`,
|
||||
);
|
||||
|
||||
// Clean up OAuth state
|
||||
cleanupOAuthState(serverName);
|
||||
|
||||
// Show error toast
|
||||
showToast({
|
||||
message: isTimeout
|
||||
? localize('com_ui_mcp_oauth_timeout', { 0: serverName })
|
||||
: localize('com_ui_mcp_init_failed'),
|
||||
status: 'error',
|
||||
});
|
||||
},
|
||||
[showToast, localize, cleanupOAuthState],
|
||||
);
|
||||
|
||||
// Poll for OAuth completion
|
||||
useEffect(() => {
|
||||
if (oauthPollingServers.size === 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
const pollInterval = setInterval(() => {
|
||||
// Check each polling server
|
||||
oauthPollingServers.forEach((oauthUrl, serverName) => {
|
||||
const serverStatus = connectionStatus[serverName];
|
||||
|
||||
// Check for client-side timeout (3 minutes)
|
||||
const startTime = oauthStartTimes.get(serverName);
|
||||
const hasTimedOut = startTime && Date.now() - startTime > 180000; // 3 minutes
|
||||
|
||||
if (serverStatus?.connectionState === 'connected') {
|
||||
// OAuth completed successfully
|
||||
handleSuccessfulConnection(
|
||||
serverName,
|
||||
localize('com_ui_mcp_authenticated_success', { 0: serverName }),
|
||||
);
|
||||
} else if (serverStatus?.connectionState === 'error' || hasTimedOut) {
|
||||
// OAuth failed or timed out
|
||||
handleOAuthFailure(serverName, !!hasTimedOut);
|
||||
}
|
||||
|
||||
setCancellableServers((prev) => new Set(prev).add(serverName));
|
||||
});
|
||||
|
||||
queryClient.refetchQueries([QueryKeys.mcpConnectionStatus]);
|
||||
}, 3500);
|
||||
|
||||
return () => {
|
||||
clearInterval(pollInterval);
|
||||
};
|
||||
}, [
|
||||
oauthPollingServers,
|
||||
oauthStartTimes,
|
||||
connectionStatus,
|
||||
queryClient,
|
||||
handleSuccessfulConnection,
|
||||
handleOAuthFailure,
|
||||
localize,
|
||||
]);
|
||||
|
||||
// Initialize server function
|
||||
const initializeServer = useCallback(
|
||||
(serverName: string) => {
|
||||
// Prevent spam - check if already initializing
|
||||
if (initializingServers.has(serverName)) {
|
||||
return;
|
||||
}
|
||||
|
||||
if (connectionStatus[serverName]?.requiresOAuth) {
|
||||
setCancellableServers((prev) => new Set(prev).add(serverName));
|
||||
}
|
||||
|
||||
// Add to initializing set
|
||||
setInitializingServers((prev) => new Set(prev).add(serverName));
|
||||
|
||||
// If there's already a server being processed, that one will be cancelled
|
||||
if (currentProcessingServer && currentProcessingServer !== serverName) {
|
||||
// Clean up the cancelled server's state immediately
|
||||
showToast({
|
||||
message: localize('com_ui_mcp_init_cancelled', { 0: currentProcessingServer }),
|
||||
status: 'warning',
|
||||
});
|
||||
|
||||
cleanupOAuthState(currentProcessingServer);
|
||||
}
|
||||
|
||||
// Track the current server being processed
|
||||
setCurrentProcessingServer(serverName);
|
||||
|
||||
reinitializeMutation.mutate(serverName, {
|
||||
onSuccess: (response: any) => {
|
||||
// Clear current processing server
|
||||
setCurrentProcessingServer(null);
|
||||
|
||||
if (response.success) {
|
||||
if (response.oauthRequired && response.oauthUrl) {
|
||||
// OAuth required - store URL and start polling
|
||||
setOauthPollingServers((prev) => new Map(prev).set(serverName, response.oauthUrl));
|
||||
|
||||
// Track when OAuth started for timeout detection
|
||||
setOauthStartTimes((prev) => new Map(prev).set(serverName, Date.now()));
|
||||
|
||||
// Call optional OAuth callback or open URL directly
|
||||
if (options?.onOAuthStarted) {
|
||||
options.onOAuthStarted(serverName, response.oauthUrl);
|
||||
} else {
|
||||
window.open(response.oauthUrl, '_blank', 'noopener,noreferrer');
|
||||
}
|
||||
|
||||
showToast({
|
||||
message: localize('com_ui_connecting'),
|
||||
status: 'info',
|
||||
});
|
||||
} else if (response.oauthRequired) {
|
||||
// OAuth required but no URL - shouldn't happen
|
||||
showToast({
|
||||
message: localize('com_ui_mcp_oauth_no_url'),
|
||||
status: 'warning',
|
||||
});
|
||||
// Remove from initializing since it failed
|
||||
setInitializingServers((prev) => {
|
||||
const newSet = new Set(prev);
|
||||
newSet.delete(serverName);
|
||||
return newSet;
|
||||
});
|
||||
} else {
|
||||
// Successful connection without OAuth
|
||||
handleSuccessfulConnection(
|
||||
serverName,
|
||||
response.message || localize('com_ui_mcp_initialized_success', { 0: serverName }),
|
||||
);
|
||||
}
|
||||
} else {
|
||||
// Remove from initializing if not successful
|
||||
setInitializingServers((prev) => {
|
||||
const newSet = new Set(prev);
|
||||
newSet.delete(serverName);
|
||||
return newSet;
|
||||
});
|
||||
}
|
||||
},
|
||||
onError: (error: any) => {
|
||||
console.error(`Error initializing MCP server ${serverName}:`, error);
|
||||
setCurrentProcessingServer(null);
|
||||
|
||||
const isCancelled =
|
||||
error?.name === 'CanceledError' ||
|
||||
error?.code === 'ERR_CANCELED' ||
|
||||
error?.message?.includes('cancel') ||
|
||||
error?.message?.includes('abort');
|
||||
|
||||
if (isCancelled) {
|
||||
showToast({
|
||||
message: localize('com_ui_mcp_init_cancelled', { 0: serverName }),
|
||||
status: 'warning',
|
||||
});
|
||||
} else {
|
||||
showToast({
|
||||
message: localize('com_ui_mcp_init_failed'),
|
||||
status: 'error',
|
||||
});
|
||||
}
|
||||
|
||||
// Clean up OAuth state using helper function
|
||||
cleanupOAuthState(serverName);
|
||||
|
||||
// Call optional error callback
|
||||
options?.onError?.(serverName, error);
|
||||
},
|
||||
});
|
||||
},
|
||||
[
|
||||
initializingServers,
|
||||
connectionStatus,
|
||||
currentProcessingServer,
|
||||
reinitializeMutation,
|
||||
showToast,
|
||||
localize,
|
||||
cleanupOAuthState,
|
||||
options,
|
||||
handleSuccessfulConnection,
|
||||
],
|
||||
);
|
||||
|
||||
return {
|
||||
initializeServer,
|
||||
isInitializing: (serverName: string) => initializingServers.has(serverName),
|
||||
isCancellable: (serverName: string) => cancellableServers.has(serverName),
|
||||
initializingServers,
|
||||
oauthPollingServers,
|
||||
oauthStartTimes,
|
||||
connectionStatus,
|
||||
isLoading: reinitializeMutation.isLoading,
|
||||
cancelOAuthFlow,
|
||||
};
|
||||
}
|
||||
|
|
@ -1,33 +1,52 @@
|
|||
import { useCallback, useState, useMemo, useRef } from 'react';
|
||||
import { useCallback, useState, useMemo, useRef, useEffect } from 'react';
|
||||
import { useToastContext } from '@librechat/client';
|
||||
import { useQueryClient } from '@tanstack/react-query';
|
||||
import { Constants, QueryKeys } from 'librechat-data-provider';
|
||||
import { useUpdateUserPluginsMutation } from 'librechat-data-provider/react-query';
|
||||
import {
|
||||
useUpdateUserPluginsMutation,
|
||||
useReinitializeMCPServerMutation,
|
||||
useCancelMCPOAuthMutation,
|
||||
} from 'librechat-data-provider/react-query';
|
||||
import { useMCPConnectionStatusQuery } from '~/data-provider/Tools/queries';
|
||||
import type { TUpdateUserPlugins, TPlugin } from 'librechat-data-provider';
|
||||
import type { ConfigFieldDetail } from '~/components/MCP/MCPConfigDialog';
|
||||
import { useLocalize, useMCPServerInitialization } from '~/hooks';
|
||||
import { useLocalize } from '~/hooks';
|
||||
import { useBadgeRowContext } from '~/Providers';
|
||||
|
||||
interface ServerState {
|
||||
isInitializing: boolean;
|
||||
oauthUrl: string | null;
|
||||
oauthStartTime: number | null;
|
||||
isCancellable: boolean;
|
||||
pollInterval: NodeJS.Timeout | null;
|
||||
}
|
||||
|
||||
export function useMCPServerManager() {
|
||||
const localize = useLocalize();
|
||||
const { showToast } = useToastContext();
|
||||
const { mcpSelect, startupConfig } = useBadgeRowContext();
|
||||
const { mcpValues, setMCPValues, mcpToolDetails, isPinned, setIsPinned } = mcpSelect;
|
||||
const queryClient = useQueryClient();
|
||||
|
||||
const [isConfigModalOpen, setIsConfigModalOpen] = useState(false);
|
||||
const [selectedToolForConfig, setSelectedToolForConfig] = useState<TPlugin | null>(null);
|
||||
const previousFocusRef = useRef<HTMLElement | null>(null);
|
||||
const mcpValuesRef = useRef(mcpValues);
|
||||
|
||||
// fixes the issue where OAuth flows would deselect all the servers except the one that is being authenticated on success
|
||||
useEffect(() => {
|
||||
mcpValuesRef.current = mcpValues;
|
||||
}, [mcpValues]);
|
||||
|
||||
const configuredServers = useMemo(() => {
|
||||
if (!startupConfig?.mcpServers) {
|
||||
return [];
|
||||
}
|
||||
if (!startupConfig?.mcpServers) return [];
|
||||
return Object.entries(startupConfig.mcpServers)
|
||||
.filter(([, config]) => config.chatMenu !== false)
|
||||
.map(([serverName]) => serverName);
|
||||
}, [startupConfig?.mcpServers]);
|
||||
|
||||
const [isConfigModalOpen, setIsConfigModalOpen] = useState(false);
|
||||
const [selectedToolForConfig, setSelectedToolForConfig] = useState<TPlugin | null>(null);
|
||||
const previousFocusRef = useRef<HTMLElement | null>(null);
|
||||
|
||||
const queryClient = useQueryClient();
|
||||
const reinitializeMutation = useReinitializeMCPServerMutation();
|
||||
const cancelOAuthMutation = useCancelMCPOAuthMutation();
|
||||
|
||||
const updateUserPluginsMutation = useUpdateUserPluginsMutation({
|
||||
onSuccess: async () => {
|
||||
|
|
@ -48,52 +67,278 @@ export function useMCPServerManager() {
|
|||
},
|
||||
});
|
||||
|
||||
const { initializeServer, isInitializing, connectionStatus, cancelOAuthFlow, isCancellable } =
|
||||
useMCPServerInitialization({
|
||||
onSuccess: (serverName) => {
|
||||
const currentValues = mcpValues ?? [];
|
||||
if (!currentValues.includes(serverName)) {
|
||||
setMCPValues([...currentValues, serverName]);
|
||||
}
|
||||
},
|
||||
onError: (serverName) => {
|
||||
const tool = mcpToolDetails?.find((t) => t.name === serverName);
|
||||
const serverConfig = startupConfig?.mcpServers?.[serverName];
|
||||
const serverStatus = connectionStatus[serverName];
|
||||
|
||||
const hasAuthConfig =
|
||||
(tool?.authConfig && tool.authConfig.length > 0) ||
|
||||
(serverConfig?.customUserVars && Object.keys(serverConfig.customUserVars).length > 0);
|
||||
|
||||
const wouldShowButton =
|
||||
!serverStatus ||
|
||||
serverStatus.connectionState === 'disconnected' ||
|
||||
serverStatus.connectionState === 'error' ||
|
||||
(serverStatus.connectionState === 'connected' && hasAuthConfig);
|
||||
|
||||
if (!wouldShowButton) {
|
||||
return;
|
||||
}
|
||||
|
||||
const configTool = tool || {
|
||||
name: serverName,
|
||||
pluginKey: `${Constants.mcp_prefix}${serverName}`,
|
||||
authConfig: serverConfig?.customUserVars
|
||||
? Object.entries(serverConfig.customUserVars).map(([key, config]) => ({
|
||||
authField: key,
|
||||
label: config.title,
|
||||
description: config.description,
|
||||
}))
|
||||
: [],
|
||||
authenticated: false,
|
||||
};
|
||||
|
||||
previousFocusRef.current = document.activeElement as HTMLElement;
|
||||
|
||||
setSelectedToolForConfig(configTool);
|
||||
setIsConfigModalOpen(true);
|
||||
},
|
||||
const [serverStates, setServerStates] = useState<Record<string, ServerState>>(() => {
|
||||
const initialStates: Record<string, ServerState> = {};
|
||||
configuredServers.forEach((serverName) => {
|
||||
initialStates[serverName] = {
|
||||
isInitializing: false,
|
||||
oauthUrl: null,
|
||||
oauthStartTime: null,
|
||||
isCancellable: false,
|
||||
pollInterval: null,
|
||||
};
|
||||
});
|
||||
return initialStates;
|
||||
});
|
||||
|
||||
const { data: connectionStatusData } = useMCPConnectionStatusQuery();
|
||||
const connectionStatus = useMemo(
|
||||
() => connectionStatusData?.connectionStatus || {},
|
||||
[connectionStatusData?.connectionStatus],
|
||||
);
|
||||
|
||||
useEffect(() => {
|
||||
if (!mcpValues?.length) return;
|
||||
|
||||
const connectedSelected = mcpValues.filter(
|
||||
(serverName) => connectionStatus[serverName]?.connectionState === 'connected',
|
||||
);
|
||||
|
||||
if (connectedSelected.length !== mcpValues.length) {
|
||||
setMCPValues(connectedSelected);
|
||||
}
|
||||
}, [connectionStatus, mcpValues, setMCPValues]);
|
||||
|
||||
const updateServerState = useCallback((serverName: string, updates: Partial<ServerState>) => {
|
||||
setServerStates((prev) => {
|
||||
const newStates = { ...prev };
|
||||
const currentState = newStates[serverName] || {
|
||||
isInitializing: false,
|
||||
oauthUrl: null,
|
||||
oauthStartTime: null,
|
||||
isCancellable: false,
|
||||
pollInterval: null,
|
||||
};
|
||||
newStates[serverName] = { ...currentState, ...updates };
|
||||
return newStates;
|
||||
});
|
||||
}, []);
|
||||
|
||||
const cleanupServerState = useCallback(
|
||||
(serverName: string) => {
|
||||
const state = serverStates[serverName];
|
||||
if (state?.pollInterval) {
|
||||
clearInterval(state.pollInterval);
|
||||
}
|
||||
updateServerState(serverName, {
|
||||
isInitializing: false,
|
||||
oauthUrl: null,
|
||||
oauthStartTime: null,
|
||||
isCancellable: false,
|
||||
pollInterval: null,
|
||||
});
|
||||
},
|
||||
[serverStates, updateServerState],
|
||||
);
|
||||
|
||||
const startServerPolling = useCallback(
|
||||
(serverName: string) => {
|
||||
const pollInterval = setInterval(async () => {
|
||||
try {
|
||||
await queryClient.refetchQueries([QueryKeys.mcpConnectionStatus]);
|
||||
|
||||
const freshConnectionData = queryClient.getQueryData([
|
||||
QueryKeys.mcpConnectionStatus,
|
||||
]) as any;
|
||||
const freshConnectionStatus = freshConnectionData?.connectionStatus || {};
|
||||
|
||||
const state = serverStates[serverName];
|
||||
const serverStatus = freshConnectionStatus[serverName];
|
||||
|
||||
if (serverStatus?.connectionState === 'connected') {
|
||||
clearInterval(pollInterval);
|
||||
|
||||
showToast({
|
||||
message: localize('com_ui_mcp_authenticated_success', { 0: serverName }),
|
||||
status: 'success',
|
||||
});
|
||||
|
||||
const currentValues = mcpValuesRef.current ?? [];
|
||||
if (!currentValues.includes(serverName)) {
|
||||
setMCPValues([...currentValues, serverName]);
|
||||
}
|
||||
|
||||
// This delay is to ensure UI has updated with new connection status before cleanup
|
||||
// Otherwise servers will show as disconnected for a second after OAuth flow completes
|
||||
setTimeout(() => {
|
||||
cleanupServerState(serverName);
|
||||
}, 1000);
|
||||
return;
|
||||
}
|
||||
|
||||
if (state?.oauthStartTime && Date.now() - state.oauthStartTime > 180000) {
|
||||
showToast({
|
||||
message: localize('com_ui_mcp_oauth_timeout', { 0: serverName }),
|
||||
status: 'error',
|
||||
});
|
||||
cleanupServerState(serverName);
|
||||
return;
|
||||
}
|
||||
|
||||
if (serverStatus?.connectionState === 'error') {
|
||||
showToast({
|
||||
message: localize('com_ui_mcp_init_failed'),
|
||||
status: 'error',
|
||||
});
|
||||
cleanupServerState(serverName);
|
||||
}
|
||||
} catch (error) {
|
||||
console.error(`[MCP Manager] Error polling server ${serverName}:`, error);
|
||||
}
|
||||
}, 3500);
|
||||
|
||||
updateServerState(serverName, { pollInterval });
|
||||
},
|
||||
[
|
||||
queryClient,
|
||||
serverStates,
|
||||
showToast,
|
||||
localize,
|
||||
setMCPValues,
|
||||
cleanupServerState,
|
||||
updateServerState,
|
||||
],
|
||||
);
|
||||
|
||||
const initializeServer = useCallback(
|
||||
async (serverName: string) => {
|
||||
updateServerState(serverName, { isInitializing: true });
|
||||
|
||||
try {
|
||||
const response = await reinitializeMutation.mutateAsync(serverName);
|
||||
|
||||
if (response.success) {
|
||||
if (response.oauthRequired && response.oauthUrl) {
|
||||
updateServerState(serverName, {
|
||||
oauthUrl: response.oauthUrl,
|
||||
oauthStartTime: Date.now(),
|
||||
isCancellable: true,
|
||||
isInitializing: true,
|
||||
});
|
||||
|
||||
window.open(response.oauthUrl, '_blank', 'noopener,noreferrer');
|
||||
|
||||
startServerPolling(serverName);
|
||||
} else {
|
||||
await queryClient.refetchQueries([QueryKeys.mcpConnectionStatus]);
|
||||
|
||||
showToast({
|
||||
message: localize('com_ui_mcp_initialized_success', { 0: serverName }),
|
||||
status: 'success',
|
||||
});
|
||||
|
||||
const currentValues = mcpValues ?? [];
|
||||
if (!currentValues.includes(serverName)) {
|
||||
setMCPValues([...currentValues, serverName]);
|
||||
}
|
||||
|
||||
cleanupServerState(serverName);
|
||||
}
|
||||
}
|
||||
} catch (error) {
|
||||
console.error(`[MCP Manager] Failed to initialize ${serverName}:`, error);
|
||||
showToast({
|
||||
message: localize('com_ui_mcp_init_failed', { 0: serverName }),
|
||||
status: 'error',
|
||||
});
|
||||
cleanupServerState(serverName);
|
||||
}
|
||||
},
|
||||
[
|
||||
updateServerState,
|
||||
reinitializeMutation,
|
||||
startServerPolling,
|
||||
queryClient,
|
||||
showToast,
|
||||
localize,
|
||||
mcpValues,
|
||||
cleanupServerState,
|
||||
setMCPValues,
|
||||
],
|
||||
);
|
||||
|
||||
const cancelOAuthFlow = useCallback(
|
||||
(serverName: string) => {
|
||||
queryClient.invalidateQueries([QueryKeys.mcpConnectionStatus]);
|
||||
cleanupServerState(serverName);
|
||||
cancelOAuthMutation.mutate(serverName);
|
||||
|
||||
showToast({
|
||||
message: localize('com_ui_mcp_oauth_cancelled', { 0: serverName }),
|
||||
status: 'warning',
|
||||
});
|
||||
},
|
||||
[queryClient, cleanupServerState, showToast, localize, cancelOAuthMutation],
|
||||
);
|
||||
|
||||
const isInitializing = useCallback(
|
||||
(serverName: string) => {
|
||||
return serverStates[serverName]?.isInitializing || false;
|
||||
},
|
||||
[serverStates],
|
||||
);
|
||||
|
||||
const isCancellable = useCallback(
|
||||
(serverName: string) => {
|
||||
return serverStates[serverName]?.isCancellable || false;
|
||||
},
|
||||
[serverStates],
|
||||
);
|
||||
|
||||
const getOAuthUrl = useCallback(
|
||||
(serverName: string) => {
|
||||
return serverStates[serverName]?.oauthUrl || null;
|
||||
},
|
||||
[serverStates],
|
||||
);
|
||||
|
||||
const placeholderText = useMemo(
|
||||
() => startupConfig?.interface?.mcpServers?.placeholder || localize('com_ui_mcp_servers'),
|
||||
[startupConfig?.interface?.mcpServers?.placeholder, localize],
|
||||
);
|
||||
|
||||
const batchToggleServers = useCallback(
|
||||
(serverNames: string[]) => {
|
||||
const connectedServers: string[] = [];
|
||||
const disconnectedServers: string[] = [];
|
||||
|
||||
serverNames.forEach((serverName) => {
|
||||
const serverStatus = connectionStatus[serverName];
|
||||
if (serverStatus?.connectionState === 'connected') {
|
||||
connectedServers.push(serverName);
|
||||
} else {
|
||||
disconnectedServers.push(serverName);
|
||||
}
|
||||
});
|
||||
|
||||
setMCPValues(connectedServers);
|
||||
|
||||
disconnectedServers.forEach((serverName) => {
|
||||
initializeServer(serverName);
|
||||
});
|
||||
},
|
||||
[connectionStatus, setMCPValues, initializeServer],
|
||||
);
|
||||
|
||||
const toggleServerSelection = useCallback(
|
||||
(serverName: string) => {
|
||||
const currentValues = mcpValues ?? [];
|
||||
const isCurrentlySelected = currentValues.includes(serverName);
|
||||
|
||||
if (isCurrentlySelected) {
|
||||
const filteredValues = currentValues.filter((name) => name !== serverName);
|
||||
setMCPValues(filteredValues);
|
||||
} else {
|
||||
const serverStatus = connectionStatus[serverName];
|
||||
if (serverStatus?.connectionState === 'connected') {
|
||||
setMCPValues([...currentValues, serverName]);
|
||||
} else {
|
||||
initializeServer(serverName);
|
||||
}
|
||||
}
|
||||
},
|
||||
[mcpValues, setMCPValues, connectionStatus, initializeServer],
|
||||
);
|
||||
|
||||
const handleConfigSave = useCallback(
|
||||
(targetName: string, authData: Record<string, string>) => {
|
||||
|
|
@ -155,48 +400,6 @@ export function useMCPServerManager() {
|
|||
}
|
||||
}, []);
|
||||
|
||||
const toggleServerSelection = useCallback(
|
||||
(serverName: string) => {
|
||||
const currentValues = mcpValues ?? [];
|
||||
const serverStatus = connectionStatus[serverName];
|
||||
|
||||
if (currentValues.includes(serverName)) {
|
||||
const filteredValues = currentValues.filter((name) => name !== serverName);
|
||||
setMCPValues(filteredValues);
|
||||
} else {
|
||||
if (serverStatus?.connectionState === 'connected') {
|
||||
setMCPValues([...currentValues, serverName]);
|
||||
} else {
|
||||
initializeServer(serverName);
|
||||
}
|
||||
}
|
||||
},
|
||||
[connectionStatus, mcpValues, setMCPValues, initializeServer],
|
||||
);
|
||||
|
||||
const batchToggleServers = useCallback(
|
||||
(serverNames: string[]) => {
|
||||
const connectedServers: string[] = [];
|
||||
const disconnectedServers: string[] = [];
|
||||
|
||||
serverNames.forEach((serverName) => {
|
||||
const serverStatus = connectionStatus[serverName];
|
||||
if (serverStatus?.connectionState === 'connected') {
|
||||
connectedServers.push(serverName);
|
||||
} else {
|
||||
disconnectedServers.push(serverName);
|
||||
}
|
||||
});
|
||||
|
||||
setMCPValues(connectedServers);
|
||||
|
||||
disconnectedServers.forEach((serverName) => {
|
||||
initializeServer(serverName);
|
||||
});
|
||||
},
|
||||
[connectionStatus, setMCPValues, initializeServer],
|
||||
);
|
||||
|
||||
const getServerStatusIconProps = useCallback(
|
||||
(serverName: string) => {
|
||||
const tool = mcpToolDetails?.find((t) => t.name === serverName);
|
||||
|
|
@ -255,11 +458,6 @@ export function useMCPServerManager() {
|
|||
],
|
||||
);
|
||||
|
||||
const placeholderText = useMemo(
|
||||
() => startupConfig?.interface?.mcpServers?.placeholder || localize('com_ui_mcp_servers'),
|
||||
[startupConfig?.interface?.mcpServers?.placeholder, localize],
|
||||
);
|
||||
|
||||
const getConfigDialogProps = useCallback(() => {
|
||||
if (!selectedToolForConfig) return null;
|
||||
|
||||
|
|
@ -302,27 +500,31 @@ export function useMCPServerManager() {
|
|||
]);
|
||||
|
||||
return {
|
||||
// Data
|
||||
configuredServers,
|
||||
connectionStatus,
|
||||
initializeServer,
|
||||
cancelOAuthFlow,
|
||||
isInitializing,
|
||||
isCancellable,
|
||||
getOAuthUrl,
|
||||
mcpValues,
|
||||
setMCPValues,
|
||||
|
||||
mcpToolDetails,
|
||||
isPinned,
|
||||
setIsPinned,
|
||||
startupConfig,
|
||||
connectionStatus,
|
||||
placeholderText,
|
||||
|
||||
// Handlers
|
||||
toggleServerSelection,
|
||||
batchToggleServers,
|
||||
getServerStatusIconProps,
|
||||
|
||||
// Dialog state
|
||||
selectedToolForConfig,
|
||||
isConfigModalOpen,
|
||||
getConfigDialogProps,
|
||||
|
||||
// Utilities
|
||||
toggleServerSelection,
|
||||
localize,
|
||||
|
||||
isConfigModalOpen,
|
||||
handleDialogOpenChange,
|
||||
selectedToolForConfig,
|
||||
setSelectedToolForConfig,
|
||||
handleSave,
|
||||
handleRevoke,
|
||||
getServerStatusIconProps,
|
||||
getConfigDialogProps,
|
||||
};
|
||||
}
|
||||
|
|
|
|||
|
|
@ -81,13 +81,21 @@ export function useMCPSelect({ conversationId }: UseMCPSelectOptions) {
|
|||
[setEphemeralAgent],
|
||||
);
|
||||
|
||||
const [mcpValues, setMCPValues] = useLocalStorage<string[]>(
|
||||
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,
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue