diff --git a/client/src/hooks/Input/mentions.spec.ts b/client/src/hooks/Input/mentions.spec.ts new file mode 100644 index 0000000000..4eefed12c5 --- /dev/null +++ b/client/src/hooks/Input/mentions.spec.ts @@ -0,0 +1,50 @@ +import { EModelEndpoint } from 'librechat-data-provider'; +import { filterMentionEndpoints } from './mentions'; + +describe('filterMentionEndpoints', () => { + const endpoints = [EModelEndpoint.anthropic, EModelEndpoint.bedrock, EModelEndpoint.agents]; + + it('limits mention endpoints to model spec addedEndpoints', () => { + const result = filterMentionEndpoints({ + endpoints, + includedEndpoints: new Set([EModelEndpoint.agents]), + includeAssistants: true, + hasAgentAccess: true, + }); + + expect(result).toEqual([EModelEndpoint.agents]); + }); + + it('keeps provider endpoints when no model spec allow-list is configured', () => { + const result = filterMentionEndpoints({ + endpoints, + includedEndpoints: new Set(), + includeAssistants: true, + hasAgentAccess: true, + }); + + expect(result).toEqual(endpoints); + }); + + it('excludes agents when the user lacks agent access', () => { + const result = filterMentionEndpoints({ + endpoints, + includedEndpoints: new Set([EModelEndpoint.agents]), + includeAssistants: true, + hasAgentAccess: false, + }); + + expect(result).toEqual([]); + }); + + it('excludes assistants when they are not included for the mention menu', () => { + const result = filterMentionEndpoints({ + endpoints: [EModelEndpoint.assistants, EModelEndpoint.azureAssistants, EModelEndpoint.openAI], + includedEndpoints: new Set(), + includeAssistants: false, + hasAgentAccess: true, + }); + + expect(result).toEqual([EModelEndpoint.openAI]); + }); +}); diff --git a/client/src/hooks/Input/mentions.ts b/client/src/hooks/Input/mentions.ts new file mode 100644 index 0000000000..7190b9ae35 --- /dev/null +++ b/client/src/hooks/Input/mentions.ts @@ -0,0 +1,31 @@ +import { EModelEndpoint, isAgentsEndpoint, isAssistantsEndpoint } from 'librechat-data-provider'; + +export function filterMentionEndpoints({ + endpoints, + includedEndpoints, + includeAssistants, + hasAgentAccess, +}: { + endpoints: Array; + includedEndpoints: Set; + includeAssistants: boolean; + hasAgentAccess: boolean; +}) { + const hasEndpointAllowList = includedEndpoints.size > 0; + + return endpoints.filter((endpoint) => { + if (!includeAssistants && isAssistantsEndpoint(endpoint)) { + return false; + } + + if (isAgentsEndpoint(endpoint) && !hasAgentAccess) { + return false; + } + + if (hasEndpointAllowList && !includedEndpoints.has(endpoint)) { + return false; + } + + return true; + }); +} diff --git a/client/src/hooks/Input/useMentions.ts b/client/src/hooks/Input/useMentions.ts index 0c16bcc1ce..420f823e23 100644 --- a/client/src/hooks/Input/useMentions.ts +++ b/client/src/hooks/Input/useMentions.ts @@ -23,6 +23,7 @@ import { useAgentsMapContext } from '~/Providers/AgentsMapContext'; import { mapEndpoints, getPresetTitle } from '~/utils'; import { EndpointIcon } from '~/components/Endpoints'; import useHasAccess from '~/hooks/Roles/useHasAccess'; +import { filterMentionEndpoints } from './mentions'; const defaultInterface = getConfigDefaults().interface; @@ -82,7 +83,25 @@ export default function useMentions({ () => startupConfig?.interface ?? defaultInterface, [startupConfig?.interface], ); - const agentQueryEnabled = hasAgentAccess && interfaceConfig.modelSelect === true; + const includedEndpoints = useMemo( + () => new Set(startupConfig?.modelSpecs?.addedEndpoints ?? []), + [startupConfig?.modelSpecs?.addedEndpoints], + ); + const validEndpoints = useMemo( + () => + filterMentionEndpoints({ + endpoints, + includedEndpoints, + includeAssistants, + hasAgentAccess, + }), + [endpoints, includedEndpoints, includeAssistants, hasAgentAccess], + ); + const validEndpointSet = useMemo(() => new Set(validEndpoints), [validEndpoints]); + const agentQueryEnabled = + hasAgentAccess && + interfaceConfig.modelSelect === true && + (includedEndpoints.size === 0 || includedEndpoints.has(EModelEndpoint.agents)); const { data: agentsList = null, isLoading: isLoadingAgents } = useListAgentsQuery( { requiredPermission: PermissionBits.VIEW }, { @@ -152,11 +171,6 @@ export default function useMentions({ }, [startupConfig, agentsMap]); const options: MentionOption[] = useMemo(() => { - let validEndpoints = endpoints; - if (!includeAssistants) { - validEndpoints = endpoints.filter((endpoint) => !isAssistantsEndpoint(endpoint)); - } - const modelOptions = validEndpoints.flatMap((endpoint) => { if (isAssistantsEndpoint(endpoint) || isAgentsEndpoint(endpoint)) { return []; @@ -207,14 +221,18 @@ export default function useMentions({ size: 20, }), })), - ...(interfaceConfig.modelSelect === true ? (agentsList ?? []) : []), + ...(interfaceConfig.modelSelect === true && validEndpointSet.has(EModelEndpoint.agents) + ? (agentsList ?? []) + : []), ...(endpointsConfig?.[EModelEndpoint.assistants] && includeAssistants && + validEndpointSet.has(EModelEndpoint.assistants) && interfaceConfig.modelSelect === true ? assistantListMap[EModelEndpoint.assistants] || [] : []), ...(endpointsConfig?.[EModelEndpoint.azureAssistants] && includeAssistants && + validEndpointSet.has(EModelEndpoint.azureAssistants) && interfaceConfig.modelSelect === true ? assistantListMap[EModelEndpoint.azureAssistants] || [] : []), @@ -241,11 +259,12 @@ export default function useMentions({ return mentions; }, [ presets, - endpoints, modelSpecs, agentsList, assistantMap, modelsConfig, + validEndpoints, + validEndpointSet, endpointsConfig, assistantListMap, includeAssistants,