feat: enhance BadgeRowContext with MCPSelect and tool toggle functionality, refactor related components to utilize updated context and hooks

client/src/Providers/BadgeRowContext.tsx

@@ -1,7 +1,12 @@
 import React, { createContext, useContext } from 'react';
+import { useMCPSelect, useToolToggle, useCodeApiKeyForm, useSearchApiKeyForm } from '~/hooks';
+import { Tools, LocalStorageKeys } from 'librechat-data-provider';
 
 interface BadgeRowContextType {
   conversationId?: string | null;
+  mcpSelect: ReturnType<typeof useMCPSelect>;
+  codeInterpreter: ReturnType<typeof useToolToggle>;
+  webSearch: ReturnType<typeof useToolToggle>;
 }
 
 const BadgeRowContext = createContext<BadgeRowContextType | undefined>(undefined);
@@ -20,8 +25,42 @@ interface BadgeRowProviderProps {
 }
 
 export default function BadgeRowProvider({ children, conversationId }: BadgeRowProviderProps) {
+  // MCPSelect hook
+  const mcpSelect = useMCPSelect({ conversationId });
+
+  // CodeInterpreter hooks
+  const { setIsDialogOpen: setCodeDialogOpen } = useCodeApiKeyForm({});
+
+  const codeInterpreter = useToolToggle({
+    conversationId,
+    setIsDialogOpen: setCodeDialogOpen,
+    toolKey: Tools.execute_code,
+    localStorageKey: LocalStorageKeys.LAST_CODE_TOGGLE_,
+    authConfig: {
+      toolId: Tools.execute_code,
+      queryOptions: { retry: 1 },
+    },
+  });
+
+  // WebSearch hooks
+  const { setIsDialogOpen: setWebSearchDialogOpen } = useSearchApiKeyForm({});
+
+  const webSearch = useToolToggle({
+    conversationId,
+    toolKey: Tools.web_search,
+    localStorageKey: LocalStorageKeys.LAST_WEB_SEARCH_TOGGLE_,
+    setIsDialogOpen: setWebSearchDialogOpen,
+    authConfig: {
+      toolId: Tools.web_search,
+      queryOptions: { retry: 1 },
+    },
+  });
+
   const value: BadgeRowContextType = {
     conversationId,
+    mcpSelect,
+    codeInterpreter,
+    webSearch,
   };
 
   return <BadgeRowContext.Provider value={value}>{children}</BadgeRowContext.Provider>;

client/src/components/Chat/Input/CodeInterpreter.tsx

@@ -1,47 +1,26 @@
 import React, { memo, useMemo, useRef } from 'react';
 import { TerminalSquareIcon } from 'lucide-react';
-import {
-  Tools,
-  AuthType,
-  PermissionTypes,
-  Permissions,
-  LocalStorageKeys,
-} from 'librechat-data-provider';
+import { AuthType, PermissionTypes, Permissions } from 'librechat-data-provider';
 import ApiKeyDialog from '~/components/SidePanel/Agents/Code/ApiKeyDialog';
-import { useLocalize, useHasAccess, useCodeApiKeyForm, useToolToggle } from '~/hooks';
+import { useLocalize, useHasAccess, useCodeApiKeyForm } from '~/hooks';
 import CheckboxButton from '~/components/ui/CheckboxButton';
-import { useVerifyAgentToolAuth } from '~/data-provider';
 import { useBadgeRowContext } from '~/Providers';
 
 function CodeInterpreter() {
   const triggerRef = useRef<HTMLInputElement>(null);
   const localize = useLocalize();
-  const { conversationId } = useBadgeRowContext();
+  const { codeInterpreter } = useBadgeRowContext();
+  const { toggleState: runCode, debouncedChange, authData } = codeInterpreter;
 
   const canRunCode = useHasAccess({
     permissionType: PermissionTypes.RUN_CODE,
     permission: Permissions.USE,
   });
 
-  const { data } = useVerifyAgentToolAuth(
-    { toolId: Tools.execute_code },
-    {
-      retry: 1,
-    },
-  );
-  const authType = useMemo(() => data?.message ?? false, [data?.message]);
-  const isAuthenticated = useMemo(() => data?.authenticated ?? false, [data?.authenticated]);
+  const authType = useMemo(() => authData?.message ?? false, [authData?.message]);
   const { methods, onSubmit, isDialogOpen, setIsDialogOpen, handleRevokeApiKey } =
     useCodeApiKeyForm({});
 
-  const { toggleState: runCode, debouncedChange } = useToolToggle({
-    conversationId,
-    isAuthenticated,
-    setIsDialogOpen,
-    toolKey: Tools.execute_code,
-    localStorageKey: LocalStorageKeys.LAST_CODE_TOGGLE_,
-  });
-
   if (!canRunCode) {
     return null;
   }
@@ -65,8 +44,8 @@ function CodeInterpreter() {
         onRevoke={handleRevokeApiKey}
         onOpenChange={setIsDialogOpen}
         handleSubmit={methods.handleSubmit}
-        isToolAuthenticated={isAuthenticated}
         isUserProvided={authType === AuthType.USER_PROVIDED}
+        isToolAuthenticated={authData?.authenticated ?? false}
       />
     </>
   );

client/src/components/Chat/Input/MCPSelect.tsx

@@ -1,23 +1,15 @@
 import React, { memo, useCallback, useState } from 'react';
 import { Settings2 } from 'lucide-react';
 import { useUpdateUserPluginsMutation } from 'librechat-data-provider/react-query';
-import { Constants, EModelEndpoint } from 'librechat-data-provider';
-import type { TPlugin, TPluginAuthConfig, TUpdateUserPlugins } from 'librechat-data-provider';
+import { Constants } from 'librechat-data-provider';
+import type { TUpdateUserPlugins } from 'librechat-data-provider';
+import type { McpServerInfo } from '~/hooks/Plugins/useMCPSelect';
 import MCPConfigDialog, { type ConfigFieldDetail } from '~/components/ui/MCPConfigDialog';
 import { useToastContext, useBadgeRowContext } from '~/Providers';
-import { useAvailableToolsQuery } from '~/data-provider';
-import { useLocalize, useMCPSelect } from '~/hooks';
 import MultiSelect from '~/components/ui/MultiSelect';
 import MCPIcon from '~/components/ui/MCPIcon';
+import { useLocalize } from '~/hooks';
 
-interface McpServerInfo {
-  name: string;
-  pluginKey: string;
-  authConfig?: TPluginAuthConfig[];
-  authenticated?: boolean;
-}
-
-// Helper function to extract mcp_serverName from a full pluginKey like action_mcp_serverName
 const getBaseMCPPluginKey = (fullPluginKey: string): string => {
   const parts = fullPluginKey.split(Constants.mcp_delimiter);
   return Constants.mcp_prefix + parts[parts.length - 1];
@@ -26,38 +18,12 @@ const getBaseMCPPluginKey = (fullPluginKey: string): string => {
 function MCPSelect() {
   const localize = useLocalize();
   const { showToast } = useToastContext();
-  const { conversationId } = useBadgeRowContext();
+  const { mcpSelect } = useBadgeRowContext();
+  const { mcpValues, setMCPValues, mcpServerNames, mcpToolDetails } = mcpSelect;
+
   const [isConfigModalOpen, setIsConfigModalOpen] = useState(false);
   const [selectedToolForConfig, setSelectedToolForConfig] = useState<McpServerInfo | null>(null);
 
-  const { data: mcpToolDetails, isFetched } = useAvailableToolsQuery(EModelEndpoint.agents, {
-    select: (data: TPlugin[]) => {
-      const mcpToolsMap = new Map<string, McpServerInfo>();
-      data.forEach((tool) => {
-        const isMCP = tool.pluginKey.includes(Constants.mcp_delimiter);
-        if (isMCP && tool.chatMenu !== false) {
-          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 { mcpValues, setMCPValues, mcpServerNames } = useMCPSelect({
-    conversationId,
-    mcpToolDetails,
-    isFetched,
-  });
-
   const updateUserPluginsMutation = useUpdateUserPluginsMutation({
     onSuccess: () => {
       setIsConfigModalOpen(false);

client/src/components/Chat/Input/WebSearch.tsx

@@ -1,41 +1,26 @@
 import React, { memo, useRef, useMemo } from 'react';
 import { Globe } from 'lucide-react';
-import { Tools, Permissions, PermissionTypes, LocalStorageKeys } from 'librechat-data-provider';
-import { useLocalize, useHasAccess, useSearchApiKeyForm, useToolToggle } from '~/hooks';
+import { Permissions, PermissionTypes } from 'librechat-data-provider';
 import ApiKeyDialog from '~/components/SidePanel/Agents/Search/ApiKeyDialog';
+import { useLocalize, useHasAccess, useSearchApiKeyForm } from '~/hooks';
 import CheckboxButton from '~/components/ui/CheckboxButton';
-import { useVerifyAgentToolAuth } from '~/data-provider';
 import { useBadgeRowContext } from '~/Providers';
 
 function WebSearch() {
   const triggerRef = useRef<HTMLInputElement>(null);
   const localize = useLocalize();
-  const { conversationId } = useBadgeRowContext();
+  const { webSearch: webSearchData } = useBadgeRowContext();
+  const { toggleState: webSearch, debouncedChange, authData } = webSearchData;
 
   const canUseWebSearch = useHasAccess({
     permissionType: PermissionTypes.WEB_SEARCH,
     permission: Permissions.USE,
   });
 
-  const { data } = useVerifyAgentToolAuth(
-    { toolId: Tools.web_search },
-    {
-      retry: 1,
-    },
-  );
-  const authTypes = useMemo(() => data?.authTypes ?? [], [data?.authTypes]);
-  const isAuthenticated = useMemo(() => data?.authenticated ?? false, [data?.authenticated]);
+  const authTypes = useMemo(() => authData?.authTypes ?? [], [authData?.authTypes]);
   const { methods, onSubmit, isDialogOpen, setIsDialogOpen, handleRevokeApiKey } =
     useSearchApiKeyForm({});
 
-  const { toggleState: webSearch, debouncedChange } = useToolToggle({
-    conversationId,
-    toolKey: Tools.web_search,
-    localStorageKey: LocalStorageKeys.LAST_WEB_SEARCH_TOGGLE_,
-    isAuthenticated,
-    setIsDialogOpen,
-  });
-
   if (!canUseWebSearch) {
     return null;
   }
@@ -60,7 +45,7 @@ function WebSearch() {
         onRevoke={handleRevokeApiKey}
         onOpenChange={setIsDialogOpen}
         handleSubmit={methods.handleSubmit}
-        isToolAuthenticated={isAuthenticated}
+        isToolAuthenticated={authData?.authenticated ?? false}
       />
     </>
   );

client/src/hooks/Plugins/useMCPSelect.ts

@@ -1,6 +1,8 @@
 import { useRef, useEffect, useCallback, useMemo } from 'react';
 import { useRecoilState } from 'recoil';
-import { Constants, LocalStorageKeys } from 'librechat-data-provider';
+import { Constants, LocalStorageKeys, EModelEndpoint } from 'librechat-data-provider';
+import type { TPlugin, TPluginAuthConfig } from 'librechat-data-provider';
+import { useAvailableToolsQuery } from '~/data-provider';
 import useLocalStorage from '~/hooks/useLocalStorageAlt';
 import { ephemeralAgentByConvoId } from '~/store';
 
@@ -20,14 +22,40 @@ const storageCondition = (value: unknown, rawCurrentValue?: string | null) => {
 
 interface UseMCPSelectOptions {
   conversationId?: string | null;
-  mcpToolDetails?: Array<{ name: string }> | null;
-  isFetched: boolean;
 }
 
-export function useMCPSelect({ conversationId, mcpToolDetails, isFetched }: UseMCPSelectOptions) {
+export interface McpServerInfo {
+  name: string;
+  pluginKey: string;
+  authConfig?: TPluginAuthConfig[];
+  authenticated?: boolean;
+}
+
+export function useMCPSelect({ conversationId }: UseMCPSelectOptions) {
   const key = conversationId ?? Constants.NEW_CONVO;
   const hasSetFetched = useRef<string | null>(null);
   const [ephemeralAgent, setEphemeralAgent] = useRecoilState(ephemeralAgentByConvoId(key));
+  const { data: mcpToolDetails, isFetched } = useAvailableToolsQuery(EModelEndpoint.agents, {
+    select: (data: TPlugin[]) => {
+      const mcpToolsMap = new Map<string, McpServerInfo>();
+      data.forEach((tool) => {
+        const isMCP = tool.pluginKey.includes(Constants.mcp_delimiter);
+        if (isMCP && tool.chatMenu !== false) {
+          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 mcpState = useMemo(() => {
     return ephemeralAgent?.mcp ?? [];
@@ -80,6 +108,7 @@ export function useMCPSelect({ conversationId, mcpToolDetails, isFetched }: UseM
     setMCPValues,
     mcpServerNames,
     ephemeralAgent,
+    mcpToolDetails,
     setEphemeralAgent,
   };
 }

client/src/hooks/Plugins/useToolToggle.ts

@@ -2,6 +2,9 @@ import { useRef, useEffect, useCallback, useMemo } from 'react';
 import { useRecoilState } from 'recoil';
 import debounce from 'lodash/debounce';
 import { Constants, LocalStorageKeys } from 'librechat-data-provider';
+import type { VerifyToolAuthResponse } from 'librechat-data-provider';
+import type { UseQueryOptions } from '@tanstack/react-query';
+import { useVerifyAgentToolAuth } from '~/data-provider';
 import useLocalStorage from '~/hooks/useLocalStorageAlt';
 import { ephemeralAgentByConvoId } from '~/store';
 
@@ -25,18 +28,35 @@ interface UseToolToggleOptions {
   localStorageKey: LocalStorageKeys;
   isAuthenticated?: boolean;
   setIsDialogOpen?: (open: boolean) => void;
+  /** Options for auth verification */
+  authConfig?: {
+    toolId: string;
+    queryOptions?: UseQueryOptions<VerifyToolAuthResponse>;
+  };
 }
 
 export function useToolToggle({
   conversationId,
   toolKey,
   localStorageKey,
-  isAuthenticated,
+  isAuthenticated: externalIsAuthenticated,
   setIsDialogOpen,
+  authConfig,
 }: UseToolToggleOptions) {
   const key = conversationId ?? Constants.NEW_CONVO;
   const [ephemeralAgent, setEphemeralAgent] = useRecoilState(ephemeralAgentByConvoId(key));
 
+  const authQuery = useVerifyAgentToolAuth(
+    { toolId: authConfig?.toolId || '' },
+    {
+      enabled: !!authConfig?.toolId,
+      ...authConfig?.queryOptions,
+    },
+  );
+
+  const isAuthenticated =
+    externalIsAuthenticated ?? (authConfig ? (authQuery?.data?.authenticated ?? false) : false);
+
   const isToolEnabled = useMemo(() => {
     return ephemeralAgent?.[toolKey] ?? false;
   }, [ephemeralAgent, toolKey]);
@@ -93,5 +113,6 @@ export function useToolToggle({
     ephemeralAgent,
     debouncedChange,
     setEphemeralAgent,
+    authData: authQuery?.data,
   };
 }