diff --git a/src/frontend/src/CustomNodes/GenericNode/components/handleRenderComponent/index.tsx b/src/frontend/src/CustomNodes/GenericNode/components/handleRenderComponent/index.tsx index 74e3fc6b4..8a1c88d7a 100644 --- a/src/frontend/src/CustomNodes/GenericNode/components/handleRenderComponent/index.tsx +++ b/src/frontend/src/CustomNodes/GenericNode/components/handleRenderComponent/index.tsx @@ -3,6 +3,7 @@ import useFlowStore from "@/stores/flowStore"; import { nodeColorsName } from "@/utils/styleUtils"; import { Connection, Handle, Position } from "@xyflow/react"; import { memo, useCallback, useEffect, useMemo, useState } from "react"; +import { useShallow } from "zustand/react/shallow"; import ShadTooltip from "../../../../components/common/shadTooltipComponent"; import { isValidConnection, @@ -179,6 +180,10 @@ const HandleRenderComponent = memo(function HandleRenderComponent({ const [isHovered, setIsHovered] = useState(false); const [openTooltip, setOpenTooltip] = useState(false); + const isLocked = useFlowStore( + useShallow((state) => state.currentFlow?.locked), + ); + const { setHandleDragging, setFilterType, @@ -384,16 +389,10 @@ const HandleRenderComponent = memo(function HandleRenderComponent({ [], ); - // Memoize the validation function - const validateConnection = useCallback( - (connection: any) => isValidConnection(connection), - [], - ); - return (
- isValidConnection(connection as Connection) + isLocked ? false : isValidConnection(connection as Connection) } className={cn( `group/handle z-50 transition-all`, !showNode && "no-show", )} - style={BASE_HANDLE_STYLES} + style={{ + ...BASE_HANDLE_STYLES, + pointerEvents: isLocked ? "none" : "auto", + }} onClick={handleClick} onMouseUp={handleMouseUp} onContextMenu={handleContextMenu} diff --git a/src/frontend/src/pages/FlowPage/components/PageComponent/index.tsx b/src/frontend/src/pages/FlowPage/components/PageComponent/index.tsx index 3d62dd09b..5fc90cc19 100644 --- a/src/frontend/src/pages/FlowPage/components/PageComponent/index.tsx +++ b/src/frontend/src/pages/FlowPage/components/PageComponent/index.tsx @@ -34,6 +34,7 @@ import { useState, } from "react"; import { useHotkeys } from "react-hotkeys-hook"; +import { useShallow } from "zustand/react/shallow"; import GenericNode from "../../../../CustomNodes/GenericNode"; import { INVALID_SELECTION_ERROR_ALERT, @@ -124,6 +125,10 @@ export default function Page({ const [selectionMenuVisible, setSelectionMenuVisible] = useState(false); const edgeUpdateSuccessful = useRef(true); + const isLocked = useFlowStore( + useShallow((state) => state.currentFlow?.locked), + ); + const position = useRef({ x: 0, y: 0 }); const [lastSelection, setLastSelection] = useState(null); @@ -521,6 +526,11 @@ export default function Page({ ); const handleEdgeClick = (event, edge) => { + if (isLocked) { + event.preventDefault(); + event.stopPropagation(); + return; + } const color = nodeColorsName[edge?.data?.sourceHandle?.output_types[0]] || "cyan"; @@ -584,13 +594,13 @@ export default function Page({ edges={edges} onNodesChange={onNodesChange} onEdgesChange={onEdgesChange} - onConnect={onConnectMod} + onConnect={isLocked ? undefined : onConnectMod} disableKeyboardA11y={true} onInit={setReactFlowInstance} nodeTypes={nodeTypes} - onReconnect={onEdgeUpdate} - onReconnectStart={onEdgeUpdateStart} - onReconnectEnd={onEdgeUpdateEnd} + onReconnect={isLocked ? undefined : onEdgeUpdate} + onReconnectStart={isLocked ? undefined : onEdgeUpdateStart} + onReconnectEnd={isLocked ? undefined : onEdgeUpdateEnd} onNodeDragStart={onNodeDragStart} onSelectionDragStart={onSelectionDragStart} elevateEdgesOnSelect={true} diff --git a/src/frontend/tests/extended/features/lock-flow.spec.ts b/src/frontend/tests/extended/features/lock-flow.spec.ts index e0435b5c7..28d6adefe 100644 --- a/src/frontend/tests/extended/features/lock-flow.spec.ts +++ b/src/frontend/tests/extended/features/lock-flow.spec.ts @@ -1,4 +1,4 @@ -import { test } from "@playwright/test"; +import { expect, Page, test } from "@playwright/test"; import * as dotenv from "dotenv"; import path from "path"; import { awaitBootstrapTest } from "../../utils/await-bootstrap-test"; @@ -74,5 +74,99 @@ test( timeout: 3000, state: "visible", }); + + await tryDeleteEdge(page); + await page.locator(".react-flow__edge-path").nth(0).click(); + await page.keyboard.press("Delete"); + let numberOfEdges = await page.locator(".react-flow__edge-path").count(); + expect(numberOfEdges).toBe(2); + + await page.locator(".react-flow__edge-path").nth(0).click(); + await page.keyboard.press("Delete"); + numberOfEdges = await page.locator(".react-flow__edge-path").count(); + expect(numberOfEdges).toBe(1); + + await page.locator(".react-flow__edge-path").nth(0).click(); + await page.keyboard.press("Delete"); + numberOfEdges = await page.locator(".react-flow__edge-path").count(); + expect(numberOfEdges).toBe(0); + + await tryConnectNodes(page); + + await page.getByTestId("handle-prompt-shownode-prompt-right").click(); + await page + .getByTestId("handle-languagemodelcomponent-shownode-system message-left") + .click(); + + await page + .getByTestId("handle-chatinput-shownode-chat message-right") + .click(); + await page + .getByTestId("handle-languagemodelcomponent-shownode-input-left") + .click(); + + await page + .getByTestId( + "handle-languagemodelcomponent-shownode-model response-right", + ) + .click(); + await page.getByTestId("handle-chatoutput-shownode-inputs-left").click(); + numberOfEdges = await page.locator(".react-flow__edge-path").count(); + + expect(numberOfEdges).toBe(3); }, ); + +async function tryConnectNodes(page: Page) { + await page.getByTestId("lock_unlock").click(); + + const numberOfTries = 5; + let numberOfEdges = await page.locator(".react-flow__edge-path").count(); + + for (let i = 0; i < numberOfTries; i++) { + try { + await page.getByTestId("handle-prompt-shownode-prompt-right").click({ + timeout: 500, + }); + } catch (e) { + numberOfEdges = await page.locator(".react-flow__edge-path").count(); + expect(numberOfEdges).toBe(0); + } + + try { + await page + .getByTestId( + "handle-languagemodelcomponent-shownode-system message-left", + ) + .click({ + timeout: 500, + }); + } catch (e) { + numberOfEdges = await page.locator(".react-flow__edge-path").count(); + expect(numberOfEdges).toBe(0); + } + } + + await page.getByTestId("lock_unlock").click(); +} + +async function tryDeleteEdge(page: Page) { + await page.getByTestId("lock_unlock").click(); + + let numberOfEdges = await page.locator(".react-flow__edge-path").count(); + expect(numberOfEdges).toBe(3); + const numberOfTries = 50; + + for (let i = 0; i < numberOfTries; i++) { + await page.locator(".react-flow__edge-path").nth(0).click(); + await page.keyboard.press("Delete"); + await page.locator(".react-flow__edge-path").nth(1).click(); + await page.keyboard.press("Delete"); + await page.locator(".react-flow__edge-path").nth(2).click(); + await page.keyboard.press("Delete"); + numberOfEdges = await page.locator(".react-flow__edge-path").count(); + expect(numberOfEdges).toBe(3); + } + //unlock the flow + await page.getByTestId("lock_unlock").click(); +}