From 01625b1b4a9a5e47068459a855288d483f3fb249 Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Sat, 21 Jun 2025 23:53:32 -0400 Subject: [PATCH] feat: enhance BadgeRowContext with MCPSelect and tool toggle functionality, refactor related components to utilize updated context and hooks --- client/src/Providers/BadgeRowContext.tsx | 39 +++++++++++++++ .../components/Chat/Input/CodeInterpreter.tsx | 33 +++---------- .../src/components/Chat/Input/MCPSelect.tsx | 48 +++---------------- .../src/components/Chat/Input/WebSearch.tsx | 27 +++-------- client/src/hooks/Plugins/useMCPSelect.ts | 37 ++++++++++++-- client/src/hooks/Plugins/useToolToggle.ts | 23 ++++++++- 6 files changed, 113 insertions(+), 94 deletions(-) diff --git a/client/src/Providers/BadgeRowContext.tsx b/client/src/Providers/BadgeRowContext.tsx index d26c490cbd..d39a631e83 100644 --- a/client/src/Providers/BadgeRowContext.tsx +++ b/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; + codeInterpreter: ReturnType; + webSearch: ReturnType; } const BadgeRowContext = createContext(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 {children}; diff --git a/client/src/components/Chat/Input/CodeInterpreter.tsx b/client/src/components/Chat/Input/CodeInterpreter.tsx index e6220c690b..1fab2d99c2 100644 --- a/client/src/components/Chat/Input/CodeInterpreter.tsx +++ b/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(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} /> ); diff --git a/client/src/components/Chat/Input/MCPSelect.tsx b/client/src/components/Chat/Input/MCPSelect.tsx index aed85d5d59..24db6b3020 100644 --- a/client/src/components/Chat/Input/MCPSelect.tsx +++ b/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(null); - const { data: mcpToolDetails, isFetched } = useAvailableToolsQuery(EModelEndpoint.agents, { - select: (data: TPlugin[]) => { - const mcpToolsMap = new Map(); - 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); diff --git a/client/src/components/Chat/Input/WebSearch.tsx b/client/src/components/Chat/Input/WebSearch.tsx index 138923fa0e..dcf3d94471 100644 --- a/client/src/components/Chat/Input/WebSearch.tsx +++ b/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(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} /> ); diff --git a/client/src/hooks/Plugins/useMCPSelect.ts b/client/src/hooks/Plugins/useMCPSelect.ts index 6845b3aa1e..a5b73b5c3c 100644 --- a/client/src/hooks/Plugins/useMCPSelect.ts +++ b/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(null); const [ephemeralAgent, setEphemeralAgent] = useRecoilState(ephemeralAgentByConvoId(key)); + const { data: mcpToolDetails, isFetched } = useAvailableToolsQuery(EModelEndpoint.agents, { + select: (data: TPlugin[]) => { + const mcpToolsMap = new Map(); + 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, }; } diff --git a/client/src/hooks/Plugins/useToolToggle.ts b/client/src/hooks/Plugins/useToolToggle.ts index ab2fd4e04f..51034f38fa 100644 --- a/client/src/hooks/Plugins/useToolToggle.ts +++ b/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; + }; } 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, }; }