diff --git a/web/app/components/workflow/hooks/use-tool-icon.ts b/web/app/components/workflow/hooks/use-tool-icon.ts index ba645736eb..8071b3a42b 100644 --- a/web/app/components/workflow/hooks/use-tool-icon.ts +++ b/web/app/components/workflow/hooks/use-tool-icon.ts @@ -15,6 +15,37 @@ import { import { CollectionType } from '@/app/components/tools/types' import { canFindTool } from '@/utils' import { useAllTriggerPlugins } from '@/service/use-triggers' +import type { PluginTriggerNodeType } from '../nodes/trigger-plugin/types' +import type { ToolNodeType } from '../nodes/tool/types' +import type { DataSourceNodeType } from '../nodes/data-source/types' +import type { TriggerWithProvider } from '../block-selector/types' + +const isTriggerPluginNode = (data: Node['data']): data is PluginTriggerNodeType => { + return data.type === BlockEnum.TriggerPlugin +} + +const isToolNode = (data: Node['data']): data is ToolNodeType => { + return data.type === BlockEnum.Tool +} + +const isDataSourceNode = (data: Node['data']): data is DataSourceNodeType => { + return data.type === BlockEnum.DataSource +} + +const findTriggerPluginIcon = ( + identifiers: (string | undefined)[], + triggers: TriggerWithProvider[] | undefined, +) => { + const targetTriggers = triggers || [] + for (const identifier of identifiers) { + if (!identifier) + continue + const matched = targetTriggers.find(trigger => trigger.id === identifier || canFindTool(trigger.id, identifier)) + if (matched?.icon) + return matched.icon + } + return undefined +} export const useToolIcon = (data?: Node['data']) => { const buildInTools = useStore(s => s.buildInTools) @@ -23,15 +54,66 @@ export const useToolIcon = (data?: Node['data']) => { const mcpTools = useStore(s => s.mcpTools) const dataSourceList = useStore(s => s.dataSourceList) const { data: triggerPlugins } = useAllTriggerPlugins() - // const a = useStore(s => s.data) const toolIcon = useMemo(() => { if (!data) return '' - if (data.type === BlockEnum.TriggerPlugin) { - const targetTools = triggerPlugins || [] - return targetTools.find(toolWithProvider => canFindTool(toolWithProvider.id, data.plugin_id))?.icon + if (isTriggerPluginNode(data)) { + const icon = findTriggerPluginIcon( + [ + data.plugin_id, + data.provider_id, + data.provider_name, + ], + triggerPlugins, + ) + return icon || '' } - if (data.type === BlockEnum.Tool) { + if (isToolNode(data)) { + // eslint-disable-next-line sonarjs/no-dead-store + let targetTools = buildInTools + if (data.provider_type === CollectionType.builtIn) + targetTools = buildInTools + else if (data.provider_type === CollectionType.custom) + targetTools = customTools + else if (data.provider_type === CollectionType.mcp) + targetTools = mcpTools + else + targetTools = workflowTools + return targetTools.find(toolWithProvider => canFindTool(toolWithProvider.id, data.provider_id))?.icon || '' + } + if (isDataSourceNode(data)) + return dataSourceList?.find(toolWithProvider => toolWithProvider.plugin_id === data.plugin_id)?.icon || '' + return '' + }, [data, dataSourceList, buildInTools, customTools, mcpTools, workflowTools, triggerPlugins]) + + return toolIcon +} + +export const useGetToolIcon = () => { + const workflowStore = useWorkflowStore() + const { data: triggerPlugins } = useAllTriggerPlugins() + + const getToolIcon = useCallback((data: Node['data']) => { + const { + buildInTools, + customTools, + workflowTools, + mcpTools, + dataSourceList, + } = workflowStore.getState() + + if (isTriggerPluginNode(data)) { + return findTriggerPluginIcon( + [ + data.plugin_id, + data.provider_id, + data.provider_name, + ], + triggerPlugins, + ) + } + + if (isToolNode(data)) { // eslint-disable-next-line sonarjs/no-dead-store let targetTools = buildInTools if (data.provider_type === CollectionType.builtIn) @@ -44,38 +126,12 @@ export const useToolIcon = (data?: Node['data']) => { targetTools = workflowTools return targetTools.find(toolWithProvider => canFindTool(toolWithProvider.id, data.provider_id))?.icon } - if (data.type === BlockEnum.DataSource) + + if (isDataSourceNode(data)) return dataSourceList?.find(toolWithProvider => toolWithProvider.plugin_id === data.plugin_id)?.icon - }, [data, dataSourceList, buildInTools, customTools, mcpTools, workflowTools, triggerPlugins]) - return toolIcon -} - -export const useGetToolIcon = () => { - const workflowStore = useWorkflowStore() - const getToolIcon = useCallback((data: Node['data']) => { - const { - buildInTools, - customTools, - workflowTools, - dataSourceList, - } = workflowStore.getState() - - if (data.type === BlockEnum.Tool) { - // eslint-disable-next-line sonarjs/no-dead-store - let targetTools = buildInTools - if (data.provider_type === CollectionType.builtIn) - targetTools = buildInTools - else if (data.provider_type === CollectionType.custom) - targetTools = customTools - else - targetTools = workflowTools - return targetTools.find(toolWithProvider => canFindTool(toolWithProvider.id, data.provider_id))?.icon - } - - if (data.type === BlockEnum.DataSource) - return dataSourceList?.find(toolWithProvider => toolWithProvider.plugin_id === data.plugin_id)?.icon - }, [workflowStore]) + return undefined + }, [workflowStore, triggerPlugins]) return getToolIcon }