diff --git a/client/src/components/Chat/Input/Files/AttachFileChat.tsx b/client/src/components/Chat/Input/Files/AttachFileChat.tsx index 079d1a997..f80b8aef8 100644 --- a/client/src/components/Chat/Input/Files/AttachFileChat.tsx +++ b/client/src/components/Chat/Input/Files/AttachFileChat.tsx @@ -39,6 +39,7 @@ function AttachFileChat({ ); diff --git a/client/src/components/Chat/Input/Files/AttachFileMenu.tsx b/client/src/components/Chat/Input/Files/AttachFileMenu.tsx index 466b7fddb..75f38b267 100644 --- a/client/src/components/Chat/Input/Files/AttachFileMenu.tsx +++ b/client/src/components/Chat/Input/Files/AttachFileMenu.tsx @@ -11,7 +11,13 @@ import { SharePointIcon, } from '@librechat/client'; import type { EndpointFileConfig } from 'librechat-data-provider'; -import { useLocalize, useGetAgentsConfig, useFileHandling, useAgentCapabilities } from '~/hooks'; +import { + useAgentToolPermissions, + useAgentCapabilities, + useGetAgentsConfig, + useFileHandling, + useLocalize, +} from '~/hooks'; import useSharePointFileHandling from '~/hooks/Files/useSharePointFileHandling'; import { SharePointPickerDialog } from '~/components/SharePoint'; import { useGetStartupConfig } from '~/data-provider'; @@ -21,11 +27,17 @@ import { cn } from '~/utils'; interface AttachFileMenuProps { conversationId: string; + agentId?: string | null; disabled?: boolean | null; endpointFileConfig?: EndpointFileConfig; } -const AttachFileMenu = ({ disabled, conversationId, endpointFileConfig }: AttachFileMenuProps) => { +const AttachFileMenu = ({ + agentId, + disabled, + conversationId, + endpointFileConfig, +}: AttachFileMenuProps) => { const localize = useLocalize(); const isUploadDisabled = disabled ?? false; const inputRef = useRef(null); @@ -52,6 +64,8 @@ const AttachFileMenu = ({ disabled, conversationId, endpointFileConfig }: Attach * */ const capabilities = useAgentCapabilities(agentsConfig?.capabilities ?? defaultAgentCapabilities); + const { fileSearchAllowedByAgent, codeAllowedByAgent } = useAgentToolPermissions(agentId); + const handleUploadClick = (isImage?: boolean) => { if (!inputRef.current) { return; @@ -86,7 +100,7 @@ const AttachFileMenu = ({ disabled, conversationId, endpointFileConfig }: Attach }); } - if (capabilities.fileSearchEnabled) { + if (capabilities.fileSearchEnabled && fileSearchAllowedByAgent) { items.push({ label: localize('com_ui_upload_file_search'), onClick: () => { @@ -101,7 +115,7 @@ const AttachFileMenu = ({ disabled, conversationId, endpointFileConfig }: Attach }); } - if (capabilities.codeEnabled) { + if (capabilities.codeEnabled && codeAllowedByAgent) { items.push({ label: localize('com_ui_upload_code_files'), onClick: () => { @@ -142,6 +156,8 @@ const AttachFileMenu = ({ disabled, conversationId, endpointFileConfig }: Attach setToolResource, setEphemeralAgent, sharePointEnabled, + codeAllowedByAgent, + fileSearchAllowedByAgent, setIsSharePointDialogOpen, ]); diff --git a/client/src/components/Chat/Input/Files/DragDropModal.tsx b/client/src/components/Chat/Input/Files/DragDropModal.tsx index c3b337b42..3263b05f1 100644 --- a/client/src/components/Chat/Input/Files/DragDropModal.tsx +++ b/client/src/components/Chat/Input/Files/DragDropModal.tsx @@ -2,7 +2,13 @@ import React, { useMemo } from 'react'; import { OGDialog, OGDialogTemplate } from '@librechat/client'; import { ImageUpIcon, FileSearch, TerminalSquareIcon, FileType2Icon } from 'lucide-react'; import { EToolResources, defaultAgentCapabilities } from 'librechat-data-provider'; -import { useLocalize, useGetAgentsConfig, useAgentCapabilities } from '~/hooks'; +import { + useAgentToolPermissions, + useAgentCapabilities, + useGetAgentsConfig, + useLocalize, +} from '~/hooks'; +import { useChatContext } from '~/Providers'; interface DragDropModalProps { onOptionSelect: (option: EToolResources | undefined) => void; @@ -26,6 +32,11 @@ const DragDropModal = ({ onOptionSelect, setShowModal, files, isVisible }: DragD * Use definition for agents endpoint for ephemeral agents * */ const capabilities = useAgentCapabilities(agentsConfig?.capabilities ?? defaultAgentCapabilities); + const { conversation } = useChatContext(); + const { fileSearchAllowedByAgent, codeAllowedByAgent } = useAgentToolPermissions( + conversation?.agent_id, + ); + const options = useMemo(() => { const _options: FileOption[] = [ { @@ -35,14 +46,14 @@ const DragDropModal = ({ onOptionSelect, setShowModal, files, isVisible }: DragD condition: files.every((file) => file.type?.startsWith('image/')), }, ]; - if (capabilities.fileSearchEnabled) { + if (capabilities.fileSearchEnabled && fileSearchAllowedByAgent) { _options.push({ label: localize('com_ui_upload_file_search'), value: EToolResources.file_search, icon: , }); } - if (capabilities.codeEnabled) { + if (capabilities.codeEnabled && codeAllowedByAgent) { _options.push({ label: localize('com_ui_upload_code_files'), value: EToolResources.execute_code, @@ -58,7 +69,7 @@ const DragDropModal = ({ onOptionSelect, setShowModal, files, isVisible }: DragD } return _options; - }, [capabilities, files, localize]); + }, [capabilities, files, localize, fileSearchAllowedByAgent, codeAllowedByAgent]); if (!isVisible) { return null; diff --git a/client/src/hooks/Agents/__tests__/useAgentToolPermissions.test.ts b/client/src/hooks/Agents/__tests__/useAgentToolPermissions.test.ts new file mode 100644 index 000000000..1d6da4ddb --- /dev/null +++ b/client/src/hooks/Agents/__tests__/useAgentToolPermissions.test.ts @@ -0,0 +1,180 @@ +import { renderHook } from '@testing-library/react'; +import { Tools } from 'librechat-data-provider'; +import useAgentToolPermissions from '../useAgentToolPermissions'; + +// Mock the dependencies +jest.mock('~/data-provider', () => ({ + useGetAgentByIdQuery: jest.fn(), +})); + +jest.mock('~/Providers', () => ({ + useAgentsMapContext: jest.fn(), +})); + +const mockUseGetAgentByIdQuery = jest.requireMock('~/data-provider').useGetAgentByIdQuery; +const mockUseAgentsMapContext = jest.requireMock('~/Providers').useAgentsMapContext; + +describe('useAgentToolPermissions', () => { + beforeEach(() => { + jest.clearAllMocks(); + }); + + describe('when no agentId is provided', () => { + it('should allow all tools for ephemeral agents', () => { + mockUseAgentsMapContext.mockReturnValue({}); + mockUseGetAgentByIdQuery.mockReturnValue({ data: undefined }); + + const { result } = renderHook(() => useAgentToolPermissions(null)); + + expect(result.current.fileSearchAllowedByAgent).toBe(true); + expect(result.current.codeAllowedByAgent).toBe(true); + expect(result.current.tools).toBeUndefined(); + }); + + it('should allow all tools when agentId is undefined', () => { + mockUseAgentsMapContext.mockReturnValue({}); + mockUseGetAgentByIdQuery.mockReturnValue({ data: undefined }); + + const { result } = renderHook(() => useAgentToolPermissions(undefined)); + + expect(result.current.fileSearchAllowedByAgent).toBe(true); + expect(result.current.codeAllowedByAgent).toBe(true); + expect(result.current.tools).toBeUndefined(); + }); + + it('should allow all tools when agentId is empty string', () => { + mockUseAgentsMapContext.mockReturnValue({}); + mockUseGetAgentByIdQuery.mockReturnValue({ data: undefined }); + + const { result } = renderHook(() => useAgentToolPermissions('')); + + expect(result.current.fileSearchAllowedByAgent).toBe(true); + expect(result.current.codeAllowedByAgent).toBe(true); + expect(result.current.tools).toBeUndefined(); + }); + }); + + describe('when agentId is provided but agent not found', () => { + it('should disallow all tools', () => { + mockUseAgentsMapContext.mockReturnValue({}); + mockUseGetAgentByIdQuery.mockReturnValue({ data: undefined }); + + const { result } = renderHook(() => useAgentToolPermissions('non-existent-agent')); + + expect(result.current.fileSearchAllowedByAgent).toBe(false); + expect(result.current.codeAllowedByAgent).toBe(false); + expect(result.current.tools).toBeUndefined(); + }); + }); + + describe('when agent is found with tools', () => { + it('should allow tools that are included in the agent tools array', () => { + const agentId = 'test-agent'; + const agent = { + id: agentId, + tools: [Tools.file_search], + }; + + mockUseAgentsMapContext.mockReturnValue({ [agentId]: agent }); + mockUseGetAgentByIdQuery.mockReturnValue({ data: undefined }); + + const { result } = renderHook(() => useAgentToolPermissions(agentId)); + + expect(result.current.fileSearchAllowedByAgent).toBe(true); + expect(result.current.codeAllowedByAgent).toBe(false); + expect(result.current.tools).toEqual([Tools.file_search]); + }); + + it('should allow both tools when both are included', () => { + const agentId = 'test-agent'; + const agent = { + id: agentId, + tools: [Tools.file_search, Tools.execute_code], + }; + + mockUseAgentsMapContext.mockReturnValue({ [agentId]: agent }); + mockUseGetAgentByIdQuery.mockReturnValue({ data: undefined }); + + const { result } = renderHook(() => useAgentToolPermissions(agentId)); + + expect(result.current.fileSearchAllowedByAgent).toBe(true); + expect(result.current.codeAllowedByAgent).toBe(true); + expect(result.current.tools).toEqual([Tools.file_search, Tools.execute_code]); + }); + + it('should use data from API query when available', () => { + const agentId = 'test-agent'; + const agentMapData = { + id: agentId, + tools: [Tools.file_search], + }; + const agentApiData = { + id: agentId, + tools: [Tools.execute_code, Tools.file_search], + }; + + mockUseAgentsMapContext.mockReturnValue({ [agentId]: agentMapData }); + mockUseGetAgentByIdQuery.mockReturnValue({ data: agentApiData }); + + const { result } = renderHook(() => useAgentToolPermissions(agentId)); + + // API data should take precedence + expect(result.current.fileSearchAllowedByAgent).toBe(true); + expect(result.current.codeAllowedByAgent).toBe(true); + expect(result.current.tools).toEqual([Tools.execute_code, Tools.file_search]); + }); + + it('should fallback to agent map data when API data is not available', () => { + const agentId = 'test-agent'; + const agentMapData = { + id: agentId, + tools: [Tools.execute_code], + }; + + mockUseAgentsMapContext.mockReturnValue({ [agentId]: agentMapData }); + mockUseGetAgentByIdQuery.mockReturnValue({ data: undefined }); + + const { result } = renderHook(() => useAgentToolPermissions(agentId)); + + expect(result.current.fileSearchAllowedByAgent).toBe(false); + expect(result.current.codeAllowedByAgent).toBe(true); + expect(result.current.tools).toEqual([Tools.execute_code]); + }); + }); + + describe('when agent has no tools', () => { + it('should disallow all tools with empty array', () => { + const agentId = 'test-agent'; + const agent = { + id: agentId, + tools: [], + }; + + mockUseAgentsMapContext.mockReturnValue({ [agentId]: agent }); + mockUseGetAgentByIdQuery.mockReturnValue({ data: undefined }); + + const { result } = renderHook(() => useAgentToolPermissions(agentId)); + + expect(result.current.fileSearchAllowedByAgent).toBe(false); + expect(result.current.codeAllowedByAgent).toBe(false); + expect(result.current.tools).toEqual([]); + }); + + it('should disallow all tools with undefined tools', () => { + const agentId = 'test-agent'; + const agent = { + id: agentId, + tools: undefined, + }; + + mockUseAgentsMapContext.mockReturnValue({ [agentId]: agent }); + mockUseGetAgentByIdQuery.mockReturnValue({ data: undefined }); + + const { result } = renderHook(() => useAgentToolPermissions(agentId)); + + expect(result.current.fileSearchAllowedByAgent).toBe(false); + expect(result.current.codeAllowedByAgent).toBe(false); + expect(result.current.tools).toBeUndefined(); + }); + }); +}); diff --git a/client/src/hooks/Agents/index.ts b/client/src/hooks/Agents/index.ts index 496d2c19a..b0df8398e 100644 --- a/client/src/hooks/Agents/index.ts +++ b/client/src/hooks/Agents/index.ts @@ -5,3 +5,4 @@ export type { ProcessedAgentCategory } from './useAgentCategories'; export { default as useAgentCapabilities } from './useAgentCapabilities'; export { default as useGetAgentsConfig } from './useGetAgentsConfig'; export { default as useAgentDefaultPermissionLevel } from './useAgentDefaultPermissionLevel'; +export { default as useAgentToolPermissions } from './useAgentToolPermissions'; diff --git a/client/src/hooks/Agents/useAgentToolPermissions.ts b/client/src/hooks/Agents/useAgentToolPermissions.ts new file mode 100644 index 000000000..747e4f5d0 --- /dev/null +++ b/client/src/hooks/Agents/useAgentToolPermissions.ts @@ -0,0 +1,65 @@ +import { useMemo } from 'react'; +import { Tools } from 'librechat-data-provider'; +import { useGetAgentByIdQuery } from '~/data-provider'; +import { useAgentsMapContext } from '~/Providers'; + +interface AgentToolPermissionsResult { + fileSearchAllowedByAgent: boolean; + codeAllowedByAgent: boolean; + tools: string[] | undefined; +} + +/** + * Hook to determine whether specific tools are allowed for a given agent. + * + * @param agentId - The ID of the agent. If null/undefined/empty, returns true for all tools (ephemeral agent behavior) + * @returns Object with boolean flags for file_search and execute_code permissions, plus the tools array + */ +export default function useAgentToolPermissions( + agentId: string | null | undefined, +): AgentToolPermissionsResult { + const agentsMap = useAgentsMapContext(); + + // Get the agent from the map if available + const selectedAgent = useMemo(() => { + return agentId != null && agentId !== '' ? agentsMap?.[agentId] : undefined; + }, [agentId, agentsMap]); + + // Query for agent data from the API + const { data: agentData } = useGetAgentByIdQuery(agentId ?? '', { + enabled: !!agentId, + }); + + // Get tools from either the API data or the agents map + const tools = useMemo( + () => + (agentData?.tools as string[] | undefined) || (selectedAgent?.tools as string[] | undefined), + [agentData?.tools, selectedAgent?.tools], + ); + + // Determine if file_search is allowed + const fileSearchAllowedByAgent = useMemo(() => { + // If no agentId, allow for ephemeral agents + if (!agentId) return true; + // If agentId exists but agent not found, disallow + if (!selectedAgent) return false; + // Check if the agent has the file_search tool + return tools?.includes(Tools.file_search) ?? false; + }, [agentId, selectedAgent, tools]); + + // Determine if execute_code is allowed + const codeAllowedByAgent = useMemo(() => { + // If no agentId, allow for ephemeral agents + if (!agentId) return true; + // If agentId exists but agent not found, disallow + if (!selectedAgent) return false; + // Check if the agent has the execute_code tool + return tools?.includes(Tools.execute_code) ?? false; + }, [agentId, selectedAgent, tools]); + + return { + fileSearchAllowedByAgent, + codeAllowedByAgent, + tools, + }; +} diff --git a/client/src/hooks/Files/useDragHelpers.ts b/client/src/hooks/Files/useDragHelpers.ts index d7ff122e7..0c923065b 100644 --- a/client/src/hooks/Files/useDragHelpers.ts +++ b/client/src/hooks/Files/useDragHelpers.ts @@ -15,6 +15,7 @@ import { import type { DropTargetMonitor } from 'react-dnd'; import type * as t from 'librechat-data-provider'; import store, { ephemeralAgentByConvoId } from '~/store'; +import { useAgentToolPermissions } from '~/hooks'; import useFileHandling from './useFileHandling'; export default function useDragHelpers() { @@ -22,6 +23,8 @@ export default function useDragHelpers() { const [showModal, setShowModal] = useState(false); const [draggedFiles, setDraggedFiles] = useState([]); const conversation = useRecoilValue(store.conversationByIndex(0)) || undefined; + const agentId = conversation?.agent_id ?? ''; + const { fileSearchAllowedByAgent, codeAllowedByAgent } = useAgentToolPermissions(agentId); const setEphemeralAgent = useSetRecoilState( ephemeralAgentByConvoId(conversation?.conversationId ?? Constants.NEW_CONVO), ); @@ -64,7 +67,18 @@ export default function useDragHelpers() { const fileSearchEnabled = capabilities.includes(AgentCapabilities.file_search) === true; const codeEnabled = capabilities.includes(AgentCapabilities.execute_code) === true; const ocrEnabled = capabilities.includes(AgentCapabilities.ocr) === true; - if (!codeEnabled && !fileSearchEnabled && !ocrEnabled) { + + /** Determine if dragged files are all images (enables the base image option) */ + const allImages = item.files.every((f) => f.type?.startsWith('image/')); + + const shouldShowModal = + allImages || + (fileSearchEnabled && fileSearchAllowedByAgent) || + (codeEnabled && codeAllowedByAgent) || + ocrEnabled; + + if (!shouldShowModal) { + // Fallback: directly handle files without showing modal handleFiles(item.files); return; }