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();
+}