Skip to content

Commit f925c83

Browse files
j-sanaaj-sanaa
committed
enhance deleteNode functionality to clean up connected inputs and remove descendants
1 parent aface38 commit f925c83

2 files changed

Lines changed: 339 additions & 18 deletions

File tree

‎packages/agentflow/src/infrastructure/store/AgentflowContext.test.tsx‎

Lines changed: 224 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -221,6 +221,230 @@ describe('AgentflowContext - deleteNode', () => {
221221
expect(result.current.state.edges).toHaveLength(1)
222222
expect(result.current.state.edges[0].id).toBe('edge-3-4')
223223
})
224+
225+
it('should clean up connected input values when node is deleted', () => {
226+
const initialFlow: FlowData = {
227+
nodes: [
228+
makeFlowNode('agent_0', {
229+
data: {
230+
id: 'agent_0',
231+
name: 'agent',
232+
label: 'Agent',
233+
outputAnchors: [{ id: 'agent_0-output-0', name: 'output', label: 'Output', type: 'Agent' }]
234+
}
235+
}),
236+
makeFlowNode('tool_0', {
237+
data: {
238+
id: 'tool_0',
239+
name: 'tool',
240+
label: 'Tool',
241+
inputs: [{ id: 'tool_0-input-agent-Agent', name: 'agent', label: 'Agent', type: 'Agent' }],
242+
inputAnchors: [{ id: 'tool_0-input-agent-Agent', name: 'agent', label: 'Agent', type: 'Agent' }],
243+
inputValues: {
244+
agent: '{{agent_0.data.instance}}',
245+
apiKey: 'sk-1234'
246+
},
247+
outputAnchors: []
248+
}
249+
})
250+
],
251+
edges: [
252+
makeEdge('agent_0', 'tool_0', {
253+
id: 'edge-agent-tool',
254+
sourceHandle: 'agent_0-output-0',
255+
targetHandle: 'tool_0-input-agent-Agent'
256+
})
257+
]
258+
}
259+
260+
const { result } = renderHook(() => useAgentflowContext(), {
261+
wrapper: createWrapper(initialFlow)
262+
})
263+
264+
// Delete agent_0
265+
act(() => {
266+
result.current.deleteNode('agent_0')
267+
})
268+
269+
// tool_0 should still exist
270+
expect(result.current.state.nodes).toHaveLength(1)
271+
const tool = result.current.state.nodes.find((n) => n.id === 'tool_0')
272+
expect(tool).toBeDefined()
273+
274+
// The agent input should be cleared, but apiKey should be preserved
275+
expect(tool?.data.inputValues?.agent).toBe('')
276+
expect(tool?.data.inputValues?.apiKey).toBe('sk-1234')
277+
278+
// Edge should be removed
279+
expect(result.current.state.edges).toHaveLength(0)
280+
})
281+
282+
it('should clean up list input values when node is deleted', () => {
283+
const initialFlow: FlowData = {
284+
nodes: [
285+
makeFlowNode('tool_0', {
286+
data: {
287+
id: 'tool_0',
288+
name: 'tool',
289+
label: 'Tool 1',
290+
outputAnchors: [{ id: 'tool_0-output-0', name: 'output', label: 'Output', type: 'Tool' }]
291+
}
292+
}),
293+
makeFlowNode('tool_1', {
294+
data: {
295+
id: 'tool_1',
296+
name: 'tool',
297+
label: 'Tool 2',
298+
outputAnchors: [{ id: 'tool_1-output-0', name: 'output', label: 'Output', type: 'Tool' }]
299+
}
300+
}),
301+
makeFlowNode('agent_0', {
302+
data: {
303+
id: 'agent_0',
304+
name: 'agent',
305+
label: 'Agent',
306+
inputs: [{ id: 'agent_0-input-tools-Tool', name: 'tools', label: 'Tools', type: 'Tool' }],
307+
inputAnchors: [{ id: 'agent_0-input-tools-Tool', name: 'tools', label: 'Tools', type: 'Tool', list: true } as any],
308+
inputValues: {
309+
tools: ['{{tool_0.data.instance}}', '{{tool_1.data.instance}}']
310+
},
311+
outputAnchors: []
312+
}
313+
})
314+
],
315+
edges: [
316+
makeEdge('tool_0', 'agent_0', {
317+
id: 'edge-tool0-agent',
318+
sourceHandle: 'tool_0-output-0',
319+
targetHandle: 'agent_0-input-tools-Tool'
320+
}),
321+
makeEdge('tool_1', 'agent_0', {
322+
id: 'edge-tool1-agent',
323+
sourceHandle: 'tool_1-output-0',
324+
targetHandle: 'agent_0-input-tools-Tool'
325+
})
326+
]
327+
}
328+
329+
const { result } = renderHook(() => useAgentflowContext(), {
330+
wrapper: createWrapper(initialFlow)
331+
})
332+
333+
// Delete tool_0
334+
act(() => {
335+
result.current.deleteNode('tool_0')
336+
})
337+
338+
// agent_0 should still exist with tool_1 reference
339+
expect(result.current.state.nodes).toHaveLength(2)
340+
const agent = result.current.state.nodes.find((n) => n.id === 'agent_0')
341+
expect(agent).toBeDefined()
342+
343+
// The tools list should only contain tool_1
344+
expect(agent?.data.inputValues?.tools).toEqual(['{{tool_1.data.instance}}'])
345+
346+
// Only one edge should remain
347+
expect(result.current.state.edges).toHaveLength(1)
348+
expect(result.current.state.edges[0].id).toBe('edge-tool1-agent')
349+
})
350+
351+
it('should delete parent node and all descendant nodes', () => {
352+
const initialFlow: FlowData = {
353+
nodes: [
354+
makeFlowNode('parent', {
355+
data: { id: 'parent', name: 'stickyNote', label: 'Parent', outputAnchors: [] }
356+
}),
357+
makeFlowNode('child-1', {
358+
parentNode: 'parent',
359+
data: { id: 'child-1', name: 'agent', label: 'Child 1', outputAnchors: [] }
360+
}),
361+
makeFlowNode('child-2', {
362+
parentNode: 'parent',
363+
data: { id: 'child-2', name: 'tool', label: 'Child 2', outputAnchors: [] }
364+
}),
365+
makeFlowNode('grandchild', {
366+
parentNode: 'child-1',
367+
data: { id: 'grandchild', name: 'agent', label: 'Grandchild', outputAnchors: [] }
368+
}),
369+
makeFlowNode('unrelated', {
370+
data: { id: 'unrelated', name: 'agent', label: 'Unrelated', outputAnchors: [] }
371+
})
372+
],
373+
edges: []
374+
}
375+
376+
const { result } = renderHook(() => useAgentflowContext(), {
377+
wrapper: createWrapper(initialFlow)
378+
})
379+
380+
// Delete parent node
381+
act(() => {
382+
result.current.deleteNode('parent')
383+
})
384+
385+
// Should only have the unrelated node left
386+
expect(result.current.state.nodes).toHaveLength(1)
387+
expect(result.current.state.nodes[0].id).toBe('unrelated')
388+
})
389+
390+
it('should clean up inputs for all descendant nodes before deletion', () => {
391+
const initialFlow: FlowData = {
392+
nodes: [
393+
makeFlowNode('parent', {
394+
data: { id: 'parent', name: 'stickyNote', label: 'Parent', outputAnchors: [] }
395+
}),
396+
makeFlowNode('child', {
397+
parentNode: 'parent',
398+
data: {
399+
id: 'child',
400+
name: 'agent',
401+
label: 'Child',
402+
outputAnchors: [{ id: 'child-output-0', name: 'output', label: 'Output', type: 'Agent' }]
403+
}
404+
}),
405+
makeFlowNode('target', {
406+
data: {
407+
id: 'target',
408+
name: 'tool',
409+
label: 'Target',
410+
inputs: [{ id: 'target-input-agent-Agent', name: 'agent', label: 'Agent', type: 'Agent' }],
411+
inputAnchors: [{ id: 'target-input-agent-Agent', name: 'agent', label: 'Agent', type: 'Agent' }],
412+
inputValues: {
413+
agent: '{{child.data.instance}}'
414+
},
415+
outputAnchors: []
416+
}
417+
})
418+
],
419+
edges: [
420+
makeEdge('child', 'target', {
421+
id: 'edge-child-target',
422+
sourceHandle: 'child-output-0',
423+
targetHandle: 'target-input-agent-Agent'
424+
})
425+
]
426+
}
427+
428+
const { result } = renderHook(() => useAgentflowContext(), {
429+
wrapper: createWrapper(initialFlow)
430+
})
431+
432+
// Delete parent (which should also delete child)
433+
act(() => {
434+
result.current.deleteNode('parent')
435+
})
436+
437+
// Should only have target node left
438+
expect(result.current.state.nodes).toHaveLength(1)
439+
const target = result.current.state.nodes[0]
440+
expect(target.id).toBe('target')
441+
442+
// The agent input should be cleared
443+
expect(target.data.inputValues?.agent).toBe('')
444+
445+
// Edge should be removed
446+
expect(result.current.state.edges).toHaveLength(0)
447+
})
224448
})
225449

226450
describe('AgentflowContext - duplicateNode', () => {

‎packages/agentflow/src/infrastructure/store/AgentflowContext.tsx‎

Lines changed: 115 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,99 @@ import { getUniqueNodeId } from '@/core/utils'
88

99
import { agentflowReducer, initialState, normalizeNodes } from './agentflowReducer'
1010

11+
// ========================================
12+
// Helper Functions
13+
// ========================================
14+
15+
/**
16+
* Check if a value is a connection string (e.g., "{{nodeId.data.instance}}")
17+
*/
18+
function isConnectionString(value: unknown): boolean {
19+
return typeof value === 'string' && value.startsWith('{{') && value.endsWith('}}')
20+
}
21+
22+
/**
23+
* Update IDs in anchor arrays to match a new node ID
24+
*/
25+
function updateAnchorIds(items: unknown, oldId: string, newId: string): void {
26+
if (!Array.isArray(items)) return
27+
for (const item of items) {
28+
if (item?.id) {
29+
item.id = item.id.replace(oldId, newId)
30+
}
31+
}
32+
}
33+
34+
/**
35+
* Clean up input values in nodes connected to a deleted node
36+
*/
37+
function deleteConnectedInput(deletedNodeId: string, nodes: FlowNode[], edges: FlowEdge[]): FlowNode[] {
38+
// Find all edges where the deleted node is the source
39+
const connectedEdges = edges.filter((edge) => edge.source === deletedNodeId)
40+
41+
return nodes.map((node) => {
42+
// Check if this node is a target of any connected edge
43+
const affectedEdge = connectedEdges.find((edge) => edge.target === node.id)
44+
if (!affectedEdge) return node
45+
46+
// Extract the input name from the target handle (format: nodeId-input-inputName-type)
47+
const targetInput = affectedEdge.targetHandle?.split('-')[2]
48+
if (!targetInput || !node.data.inputValues) return node
49+
50+
// Clean up the input value
51+
const currentValue = node.data.inputValues[targetInput]
52+
let newValue: unknown
53+
54+
// Check if this is a list input (array of connections)
55+
const inputAnchor = node.data.inputAnchors?.find((anchor) => anchor.name === targetInput)
56+
const inputParam = node.data.inputs?.find((param) => param.name === targetInput)
57+
58+
if (inputAnchor && (inputAnchor as { list?: boolean }).list) {
59+
// For list inputs, filter out connections to the deleted node
60+
const values = (currentValue as string[]) || []
61+
newValue = values.filter((item) => !item.includes(deletedNodeId))
62+
} else if (inputParam && (inputParam as { acceptVariable?: boolean }).acceptVariable) {
63+
// For variable inputs, remove the variable reference
64+
newValue = ((currentValue as string) || '').replace(`{{${deletedNodeId}.data.instance}}`, '') || ''
65+
} else {
66+
// Default: clear the value
67+
newValue = ''
68+
}
69+
70+
return {
71+
...node,
72+
data: {
73+
...node.data,
74+
inputValues: {
75+
...node.data.inputValues,
76+
[targetInput]: newValue
77+
}
78+
}
79+
}
80+
})
81+
}
82+
83+
/**
84+
* Recursively collect all descendant nodes of a parent
85+
*/
86+
function collectDescendants(parentId: string, nodes: FlowNode[]): Set<string> {
87+
const nodesToDelete = new Set<string>()
88+
const childNodes = nodes.filter((node) => node.parentNode === parentId)
89+
90+
childNodes.forEach((childNode) => {
91+
nodesToDelete.add(childNode.id)
92+
// Recursively collect descendants of this child
93+
const descendants = collectDescendants(childNode.id, nodes)
94+
descendants.forEach((id) => nodesToDelete.add(id))
95+
})
96+
97+
return nodesToDelete
98+
}
99+
100+
// ========================================
101+
// Types
102+
// ========================================
103+
11104
// Local state setter types
12105
type NodesSetter = (nodes: FlowNode[]) => void
13106
type EdgesSetter = (edges: FlowEdge[]) => void
@@ -123,8 +216,25 @@ export function AgentflowStateProvider({ children, initialFlow }: AgentflowState
123216
// Node operations
124217
const deleteNode = useCallback(
125218
(nodeId: string) => {
126-
const newNodes = state.nodes.filter((node) => node.id !== nodeId)
127-
const newEdges = state.edges.filter((edge) => edge.source !== nodeId && edge.target !== nodeId)
219+
// Collect all nodes to be deleted (parent and all descendants)
220+
const nodesToDelete = new Set<string>()
221+
nodesToDelete.add(nodeId)
222+
const descendants = collectDescendants(nodeId, state.nodes)
223+
descendants.forEach((id) => nodesToDelete.add(id))
224+
225+
// Clean up connected inputs for the parent node first
226+
let updatedNodes = deleteConnectedInput(nodeId, state.nodes, state.edges)
227+
228+
// Clean up connected inputs for each descendant
229+
descendants.forEach((id) => {
230+
updatedNodes = deleteConnectedInput(id, updatedNodes, state.edges)
231+
})
232+
233+
// Remove all nodes in the deletion set
234+
const newNodes = updatedNodes.filter((node) => !nodesToDelete.has(node.id))
235+
236+
// Remove all edges connected to any deleted node
237+
const newEdges = state.edges.filter((edge) => !nodesToDelete.has(edge.source) && !nodesToDelete.has(edge.target))
128238

129239
syncStateUpdate({ nodes: newNodes, edges: newEdges })
130240
},
@@ -157,23 +267,10 @@ export function AgentflowStateProvider({ children, initialFlow }: AgentflowState
157267
selected: false
158268
}
159269

160-
// Helper to update IDs in anchor arrays
161-
const updateAnchorIds = (items: unknown) => {
162-
if (!Array.isArray(items)) return
163-
for (const item of items) {
164-
if (item?.id) {
165-
item.id = item.id.replace(nodeId, newNodeId)
166-
}
167-
}
168-
}
169-
170270
// Update IDs in all anchor arrays to match new node ID
171-
updateAnchorIds(newNode.data.inputs)
172-
updateAnchorIds(newNode.data.inputAnchors)
173-
updateAnchorIds(newNode.data.outputAnchors)
174-
175-
// Helper to check if value is a connection reference
176-
const isConnectionString = (v: unknown): boolean => typeof v === 'string' && v.startsWith('{{') && v.endsWith('}}')
271+
updateAnchorIds(newNode.data.inputs, nodeId, newNodeId)
272+
updateAnchorIds(newNode.data.inputAnchors, nodeId, newNodeId)
273+
updateAnchorIds(newNode.data.outputAnchors, nodeId, newNodeId)
177274

178275
// Clear connected input values by resetting to defaults
179276
if (newNode.data.inputValues) {

0 commit comments

Comments
 (0)