From f9966c2ee4c7c4a4bc3680155085f9cf521a3ed3 Mon Sep 17 00:00:00 2001 From: KVOJJJin Date: Fri, 31 Jul 2026 14:57:10 +0800 Subject: [PATCH] fix(workflow): remap snippet references (#39843) --- .../__tests__/use-insert-snippet.spec.tsx | 126 +++++++++++++++++- .../snippets/use-insert-snippet.ts | 77 ++++++++++- 2 files changed, 199 insertions(+), 4 deletions(-) diff --git a/web/app/components/workflow/block-selector/snippets/__tests__/use-insert-snippet.spec.tsx b/web/app/components/workflow/block-selector/snippets/__tests__/use-insert-snippet.spec.tsx index 9ec90476499..6c2e2ccc1d5 100644 --- a/web/app/components/workflow/block-selector/snippets/__tests__/use-insert-snippet.spec.tsx +++ b/web/app/components/workflow/block-selector/snippets/__tests__/use-insert-snippet.spec.tsx @@ -14,6 +14,12 @@ type TestNode = { _children?: { nodeId: string; nodeType: string }[] _connectedSourceHandleIds?: string[] _connectedTargetHandleIds?: string[] + variables?: { variable: string; value_selector: string[] }[] + iteration_id?: string + loop_id?: string + start_node_id?: string + iterator_selector?: string[] + output_selector?: string[] } } @@ -112,6 +118,9 @@ describe('useInsertSnippet', () => { type: 'iteration', selected: false, _children: [{ nodeId: 'snippet-node-2', nodeType: 'code' }], + start_node_id: 'snippet-node-2', + iterator_selector: ['start', 'items'], + output_selector: ['snippet-node-2', 'result'], }, }, { @@ -119,7 +128,11 @@ describe('useInsertSnippet', () => { parentId: 'snippet-node-1', zIndex: 1002, position: { x: 30, y: 40 }, - data: { type: 'code', selected: false }, + data: { + type: 'code', + selected: false, + iteration_id: 'snippet-node-1', + }, }, ], edges: [ @@ -155,6 +168,10 @@ describe('useInsertSnippet', () => { expect(nextNodes[1]!.zIndex).toBe(0) expect(nextNodes[2]!.zIndex).toBe(NESTED_ELEMENT_Z_INDEX) expect(nextNodes[1]!.data._children![0]!.nodeId).toBe(nextNodes[2]!.id) + expect(nextNodes[1]!.data.start_node_id).toBe(nextNodes[2]!.id) + expect(nextNodes[1]!.data.iterator_selector).toEqual(['start', 'items']) + expect(nextNodes[1]!.data.output_selector).toEqual([nextNodes[2]!.id, 'result']) + expect(nextNodes[2]!.data.iteration_id).toBe(nextNodes[1]!.id) const nextEdges = mockSetEdges.mock.calls[0]![0] as TestEdge[] expect(nextEdges).toHaveLength(2) @@ -171,6 +188,106 @@ describe('useInsertSnippet', () => { }) }) + it('should remap variable selectors that reference nodes inside the snippet', async () => { + mockFetchQuery.mockResolvedValue({ + graph: { + nodes: [ + { + id: 'snippet-llm', + position: { x: 10, y: 20 }, + data: { + type: 'llm', + selected: false, + model: { mode: 'chat' }, + prompt_template: [], + }, + }, + { + id: 'snippet-code', + position: { x: 310, y: 20 }, + data: { + type: 'code', + selected: false, + variables: [ + { variable: 'arg1', value_selector: ['snippet-llm', 'text'] }, + { variable: 'arg2', value_selector: ['snippet-llm', 'text'] }, + ], + }, + }, + ], + edges: [ + { + id: 'snippet-llm-source-snippet-code-target', + source: 'snippet-llm', + sourceHandle: 'source', + target: 'snippet-code', + targetHandle: 'target', + data: {}, + }, + ], + }, + }) + + const { result } = renderHook(() => useInsertSnippet()) + + await act(async () => { + await result.current.handleInsertSnippet('snippet-1') + }) + + const nextNodes = mockSetNodes.mock.calls[0]![0] as TestNode[] + const insertedLLMNode = nextNodes.find((node) => node.id.includes('snippet-llm'))! + const insertedCodeNode = nextNodes.find((node) => node.id.includes('snippet-code'))! + + expect(insertedLLMNode.id).not.toBe('snippet-llm') + expect(insertedCodeNode.data.variables).toEqual([ + { variable: 'arg1', value_selector: [insertedLLMNode.id, 'text'] }, + { variable: 'arg2', value_selector: [insertedLLMNode.id, 'text'] }, + ]) + }) + + it('should remap structural node ids inside a loop snippet', async () => { + mockFetchQuery.mockResolvedValue({ + graph: { + nodes: [ + { + id: 'snippet-loop', + position: { x: 10, y: 20 }, + data: { + type: 'loop', + selected: false, + _children: [{ nodeId: 'snippet-loop-start', nodeType: 'loop-start' }], + start_node_id: 'snippet-loop-start', + }, + }, + { + id: 'snippet-loop-start', + parentId: 'snippet-loop', + position: { x: 30, y: 40 }, + data: { + type: 'loop-start', + selected: false, + loop_id: 'snippet-loop', + }, + }, + ], + edges: [], + }, + }) + + const { result } = renderHook(() => useInsertSnippet()) + + await act(async () => { + await result.current.handleInsertSnippet('snippet-1') + }) + + const nextNodes = mockSetNodes.mock.calls[0]![0] as TestNode[] + const insertedLoopNode = nextNodes.find((node) => node.data.type === 'loop')! + const insertedLoopStartNode = nextNodes.find((node) => node.data.type === 'loop-start')! + + expect(insertedLoopNode.data.start_node_id).toBe(insertedLoopStartNode.id) + expect(insertedLoopStartNode.data.loop_id).toBe(insertedLoopNode.id) + }) + it.each(['iteration', 'loop'] as const)( 'should connect inserted snippet nodes inside the %s container', async (containerType) => { @@ -220,7 +337,12 @@ describe('useInsertSnippet', () => { { id: 'snippet-entry', position: { x: 0, y: 0 }, - data: { type: 'llm', selected: false }, + data: { + type: 'llm', + selected: false, + model: { mode: 'chat' }, + prompt_template: [], + }, }, { id: 'snippet-exit', diff --git a/web/app/components/workflow/block-selector/snippets/use-insert-snippet.ts b/web/app/components/workflow/block-selector/snippets/use-insert-snippet.ts index 0e2e5b5de98..87fca3d5d7c 100644 --- a/web/app/components/workflow/block-selector/snippets/use-insert-snippet.ts +++ b/web/app/components/workflow/block-selector/snippets/use-insert-snippet.ts @@ -1,4 +1,6 @@ -import type { Edge, Node, OnNodeAdd } from '../../types' +import type { IterationNodeType } from '../../nodes/iteration/types' +import type { LoopNodeType } from '../../nodes/loop/types' +import type { Edge, Node, OnNodeAdd, ValueSelector } from '../../types' import { toast } from '@langgenius/dify-ui/toast' import { useQueryClient } from '@tanstack/react-query' import { useCallback } from 'react' @@ -9,6 +11,7 @@ import { useIncrementSnippetUseCountMutation } from '@/service/use-snippets' import { CUSTOM_EDGE, NESTED_ELEMENT_Z_INDEX, NODE_WIDTH_X_OFFSET, X_OFFSET } from '../../constants' import { useNodesSyncDraft } from '../../hooks/use-nodes-sync-draft' import { useWorkflowHistory, WorkflowHistoryEvent } from '../../hooks/use-workflow-history' +import { getNodeUsedVars, updateNodeVars } from '../../nodes/_base/components/variable/utils' import { BlockEnum } from '../../types' import { getNodesConnectedSourceOrTargetHandleIdsMap } from '../../utils' @@ -83,6 +86,72 @@ const getInsertAnchor = (currentNodes: Node[], insertPayload?: SnippetInsertPayl } } +const remapSnippetNodeVariableReferences = (node: Node, idMapping: Map) => { + return getNodeUsedVars(node).reduce((currentNode, selector) => { + if (!Array.isArray(selector) || !selector.length) return currentNode + + const mappedNodeId = idMapping.get(selector[0]!) + if (!mappedNodeId) return currentNode + + return updateNodeVars(currentNode, selector, [mappedNodeId, ...selector.slice(1)]) + }, node) +} + +const remapSnippetNodeId = (nodeId: string | undefined, idMapping: Map) => { + if (!nodeId) return nodeId + + return idMapping.get(nodeId) ?? nodeId +} + +const remapSnippetValueSelector = ( + selector: ValueSelector | undefined, + idMapping: Map, +) => { + if (!Array.isArray(selector) || !selector.length) return selector + + const mappedNodeId = idMapping.get(selector[0]!) + if (!mappedNodeId) return selector + + return [mappedNodeId, ...selector.slice(1)] +} + +const remapSnippetNodeStructuralReferences = (node: Node, idMapping: Map): Node => { + const data = { + ...node.data, + iteration_id: remapSnippetNodeId(node.data.iteration_id, idMapping), + loop_id: remapSnippetNodeId(node.data.loop_id, idMapping), + } + + if (data.type === BlockEnum.Iteration) { + const iterationData = data as IterationNodeType + const remappedIterationData: IterationNodeType = { + ...iterationData, + start_node_id: remapSnippetNodeId(iterationData.start_node_id, idMapping)!, + output_selector: remapSnippetValueSelector(iterationData.output_selector, idMapping)!, + } + + return { + ...node, + data: remappedIterationData, + } + } + + if (data.type === BlockEnum.Loop) { + const loopData = data as LoopNodeType + const remappedLoopData: LoopNodeType = { + ...loopData, + start_node_id: remapSnippetNodeId(loopData.start_node_id, idMapping)!, + } + + return { + ...node, + data: remappedLoopData, + } + } + + return { ...node, data } +} + const remapSnippetGraph = ( currentNodes: Node[], snippetNodes: Node[], @@ -120,7 +189,7 @@ const remapSnippetGraph = ( const nextParentId = node.parentId ? idMapping.get(node.parentId) : undefined const isRootNode = !node.parentId - return { + const remappedNode = { ...node, id: idMapping.get(node.id)!, parentId: nextParentId, @@ -148,6 +217,10 @@ const remapSnippetGraph = ( })), }, } + + const structurallyRemappedNode = remapSnippetNodeStructuralReferences(remappedNode, idMapping) + + return remapSnippetNodeVariableReferences(structurallyRemappedNode, idMapping) }) const edges = snippetEdges.map((edge) => ({