diff --git a/web/app/components/workflow/hooks/__tests__/use-available-blocks.spec.ts b/web/app/components/workflow/hooks/__tests__/use-available-blocks.spec.ts index c89ba9ce96..b8dc614fdd 100644 --- a/web/app/components/workflow/hooks/__tests__/use-available-blocks.spec.ts +++ b/web/app/components/workflow/hooks/__tests__/use-available-blocks.spec.ts @@ -81,9 +81,10 @@ describe('useAvailableBlocks', () => { expect(result.current.availableNextBlocks).toEqual([]) }) - it('should return empty array for End node', () => { + it('should return available blocks for End node (same as regular nodes)', () => { const { result } = renderWorkflowHook(() => useAvailableBlocks(BlockEnum.End), { hooksStoreProps }) - expect(result.current.availableNextBlocks).toEqual([]) + expect(result.current.availableNextBlocks.length).toBeGreaterThan(0) + expect(result.current.availableNextBlocks).toContain(BlockEnum.LLM) }) it('should return empty array for LoopEnd node', () => { @@ -143,14 +144,21 @@ describe('useAvailableBlocks', () => { expect(blocks.availablePrevBlocks).toEqual([]) }) - it('should return empty nextBlocks for End/LoopEnd/KnowledgeBase', () => { + it('should return empty nextBlocks for LoopEnd/KnowledgeBase', () => { const { result } = renderWorkflowHook(() => useAvailableBlocks(BlockEnum.LLM), { hooksStoreProps }) - expect(result.current.getAvailableBlocks(BlockEnum.End).availableNextBlocks).toEqual([]) expect(result.current.getAvailableBlocks(BlockEnum.LoopEnd).availableNextBlocks).toEqual([]) expect(result.current.getAvailableBlocks(BlockEnum.KnowledgeBase).availableNextBlocks).toEqual([]) }) + it('should return available nextBlocks for End node (same as regular nodes)', () => { + const { result } = renderWorkflowHook(() => useAvailableBlocks(BlockEnum.LLM), { hooksStoreProps }) + const blocks = result.current.getAvailableBlocks(BlockEnum.End) + + expect(blocks.availableNextBlocks.length).toBeGreaterThan(0) + expect(blocks.availableNextBlocks).toContain(BlockEnum.LLM) + }) + it('should filter by inContainer when provided', () => { const { result } = renderWorkflowHook(() => useAvailableBlocks(BlockEnum.LLM), { hooksStoreProps }) const blocks = result.current.getAvailableBlocks(BlockEnum.Code, true) diff --git a/web/app/components/workflow/hooks/use-available-blocks.ts b/web/app/components/workflow/hooks/use-available-blocks.ts index 3d92142b10..888e126474 100644 --- a/web/app/components/workflow/hooks/use-available-blocks.ts +++ b/web/app/components/workflow/hooks/use-available-blocks.ts @@ -30,7 +30,7 @@ export const useAvailableBlocks = (nodeType?: BlockEnum, inContainer?: boolean) return availableNodesType }, [availableNodesType, nodeType]) const availableNextBlocks = useMemo(() => { - if (!nodeType || nodeType === BlockEnum.End || nodeType === BlockEnum.LoopEnd || nodeType === BlockEnum.KnowledgeBase) + if (!nodeType || nodeType === BlockEnum.LoopEnd || nodeType === BlockEnum.KnowledgeBase) return [] return availableNodesType @@ -42,7 +42,7 @@ export const useAvailableBlocks = (nodeType?: BlockEnum, inContainer?: boolean) availablePrevBlocks = [] let availableNextBlocks = availableNodesType - if (!nodeType || nodeType === BlockEnum.End || nodeType === BlockEnum.LoopEnd || nodeType === BlockEnum.KnowledgeBase) + if (!nodeType || nodeType === BlockEnum.LoopEnd || nodeType === BlockEnum.KnowledgeBase) availableNextBlocks = [] return {