diff --git a/src/frontend/src/components/core/logCanvasControlsComponent/index.tsx b/src/frontend/src/components/core/logCanvasControlsComponent/index.tsx
index 2bd7fe569a29..0d3f182f4d22 100644
--- a/src/frontend/src/components/core/logCanvasControlsComponent/index.tsx
+++ b/src/frontend/src/components/core/logCanvasControlsComponent/index.tsx
@@ -7,7 +7,7 @@ const LogCanvasControls = () => {
   return (
     <Panel
       data-testid="canvas_controls"
-      className="react-flow__controls !m-2 rounded-md !left-auto "
+      className="react-flow__controls !m-2 rounded-md"
       position="bottom-left"
     >
       <FlowLogsModal>
diff --git a/src/frontend/src/pages/FlowPage/components/PageComponent/MemoizedComponents.tsx b/src/frontend/src/pages/FlowPage/components/PageComponent/MemoizedComponents.tsx
index 2d71c2a97aeb..60843cc62bd5 100644
--- a/src/frontend/src/pages/FlowPage/components/PageComponent/MemoizedComponents.tsx
+++ b/src/frontend/src/pages/FlowPage/components/PageComponent/MemoizedComponents.tsx
@@ -17,7 +17,6 @@ export const MemoizedBackground = memo(() => (
 
 interface MemoizedCanvasControlsProps {
   setIsAddingNote: (value: boolean) => void;
-  position: { x: number; y: number };
   shadowBoxWidth: number;
   shadowBoxHeight: number;
 }
@@ -27,7 +26,6 @@ export const MemoizedLogCanvasControls = memo(() => <LogCanvasControls />);
 export const MemoizedCanvasControls = memo(
   ({
     setIsAddingNote,
-    position,
     shadowBoxWidth,
     shadowBoxHeight,
   }: MemoizedCanvasControlsProps) => (
@@ -43,9 +41,9 @@ export const MemoizedCanvasControls = memo(
           setIsAddingNote(true);
           const shadowBox = document.getElementById("shadow-box");
           if (shadowBox) {
+            shadowBox.style.left = `${e.clientX - shadowBoxWidth / 2}px`;
+            shadowBox.style.top = `${e.clientY - shadowBoxHeight / 2}px`;
             shadowBox.style.display = "block";
-            shadowBox.style.left = `${position.x - shadowBoxWidth / 2}px`;
-            shadowBox.style.top = `${position.y - shadowBoxHeight / 2}px`;
           }
         }}
       >
diff --git a/src/frontend/src/pages/FlowPage/components/PageComponent/index.tsx b/src/frontend/src/pages/FlowPage/components/PageComponent/index.tsx
index 7ded70e0045a..aae515f6faec 100644
--- a/src/frontend/src/pages/FlowPage/components/PageComponent/index.tsx
+++ b/src/frontend/src/pages/FlowPage/components/PageComponent/index.tsx
@@ -676,74 +676,75 @@ export default function Page({
   return (
     <div className="h-full w-full bg-canvas" ref={reactFlowWrapper}>
       {showCanvas ? (
-        <div id="react-flow-id" className="h-full w-full bg-canvas">
-          {!view && (
-            <>
-              <MemoizedLogCanvasControls />
-              <MemoizedCanvasControls
-                setIsAddingNote={setIsAddingNote}
-                position={position.current}
-                shadowBoxWidth={shadowBoxWidth}
-                shadowBoxHeight={shadowBoxHeight}
-              />
-              <FlowToolbar />
-            </>
-          )}
-          <MemoizedSidebarTrigger />
-          <SelectionMenu
-            lastSelection={lastSelection}
-            isVisible={selectionMenuVisible}
-            nodes={lastSelection?.nodes}
-            onClick={handleGroupNode}
-          />
-          <ReactFlow<AllNodeType, EdgeType>
-            nodes={nodes}
-            edges={edges}
-            onNodesChange={onNodesChangeWithHelperLines}
-            onEdgesChange={onEdgesChange}
-            onConnect={isLocked ? undefined : onConnectMod}
-            disableKeyboardA11y={true}
-            nodesFocusable={!isLocked}
-            edgesFocusable={!isLocked}
-            onInit={setReactFlowInstance}
-            nodeTypes={nodeTypes}
-            onReconnect={isLocked ? undefined : onEdgeUpdate}
-            onReconnectStart={isLocked ? undefined : onEdgeUpdateStart}
-            onReconnectEnd={isLocked ? undefined : onEdgeUpdateEnd}
-            onNodeDrag={onNodeDrag}
-            onNodeDragStart={onNodeDragStart}
-            onSelectionDragStart={onSelectionDragStart}
-            elevateEdgesOnSelect={false}
-            onSelectionEnd={onSelectionEnd}
-            onSelectionStart={onSelectionStart}
-            connectionRadius={30}
-            edgeTypes={edgeTypes}
-            connectionLineComponent={ConnectionLineComponent}
-            onDragOver={onDragOver}
-            onNodeDragStop={onNodeDragStop}
-            onDrop={onDrop}
-            onSelectionChange={onSelectionChange}
-            deleteKeyCode={[]}
-            fitView={isEmptyFlow.current ? false : true}
-            fitViewOptions={fitViewOptions}
-            className="theme-attribution"
-            tabIndex={isLocked ? -1 : undefined}
-            minZoom={MIN_ZOOM}
-            maxZoom={MAX_ZOOM}
-            zoomOnScroll={!view}
-            zoomOnPinch={!view}
-            panOnDrag={!view}
-            panActivationKeyCode={""}
-            proOptions={{ hideAttribution: true }}
-            onPaneClick={onPaneClick}
-            onEdgeClick={handleEdgeClick}
-            onKeyDown={handleKeyDown}
-          >
-            <FlowBuildingComponent />
-            <UpdateAllComponents />
-            <MemoizedBackground />
-            {helperLineEnabled && <HelperLines helperLines={helperLines} />}
-          </ReactFlow>
+        <>
+          <div id="react-flow-id" className="h-full w-full bg-canvas relative">
+            {!view && (
+              <>
+                <MemoizedLogCanvasControls />
+                <MemoizedCanvasControls
+                  setIsAddingNote={setIsAddingNote}
+                  shadowBoxWidth={shadowBoxWidth}
+                  shadowBoxHeight={shadowBoxHeight}
+                />
+                <FlowToolbar />
+              </>
+            )}
+            <MemoizedSidebarTrigger />
+            <SelectionMenu
+              lastSelection={lastSelection}
+              isVisible={selectionMenuVisible}
+              nodes={lastSelection?.nodes}
+              onClick={handleGroupNode}
+            />
+            <ReactFlow<AllNodeType, EdgeType>
+              nodes={nodes}
+              edges={edges}
+              onNodesChange={onNodesChangeWithHelperLines}
+              onEdgesChange={onEdgesChange}
+              onConnect={isLocked ? undefined : onConnectMod}
+              disableKeyboardA11y={true}
+              nodesFocusable={!isLocked}
+              edgesFocusable={!isLocked}
+              onInit={setReactFlowInstance}
+              nodeTypes={nodeTypes}
+              onReconnect={isLocked ? undefined : onEdgeUpdate}
+              onReconnectStart={isLocked ? undefined : onEdgeUpdateStart}
+              onReconnectEnd={isLocked ? undefined : onEdgeUpdateEnd}
+              onNodeDrag={onNodeDrag}
+              onNodeDragStart={onNodeDragStart}
+              onSelectionDragStart={onSelectionDragStart}
+              elevateEdgesOnSelect={false}
+              onSelectionEnd={onSelectionEnd}
+              onSelectionStart={onSelectionStart}
+              connectionRadius={30}
+              edgeTypes={edgeTypes}
+              connectionLineComponent={ConnectionLineComponent}
+              onDragOver={onDragOver}
+              onNodeDragStop={onNodeDragStop}
+              onDrop={onDrop}
+              onSelectionChange={onSelectionChange}
+              deleteKeyCode={[]}
+              fitView={isEmptyFlow.current ? false : true}
+              fitViewOptions={fitViewOptions}
+              className="theme-attribution"
+              tabIndex={isLocked ? -1 : undefined}
+              minZoom={MIN_ZOOM}
+              maxZoom={MAX_ZOOM}
+              zoomOnScroll={!view}
+              zoomOnPinch={!view}
+              panOnDrag={!view}
+              panActivationKeyCode={""}
+              proOptions={{ hideAttribution: true }}
+              onPaneClick={onPaneClick}
+              onEdgeClick={handleEdgeClick}
+              onKeyDown={handleKeyDown}
+            >
+              <FlowBuildingComponent />
+              <UpdateAllComponents />
+              <MemoizedBackground />
+              {helperLineEnabled && <HelperLines helperLines={helperLines} />}
+            </ReactFlow>
+          </div>
           <div
             id="shadow-box"
             style={{
@@ -757,7 +758,7 @@ export default function Page({
               display: "none",
             }}
           ></div>
-        </div>
+        </>
       ) : (
         <div className="flex h-full w-full items-center justify-center">
           <CustomLoader remSize={30} />
diff --git a/src/frontend/tests/extended/features/lock-flow.spec.ts b/src/frontend/tests/extended/features/lock-flow.spec.ts
index ecb1ff9f40ea..479a8ee5b681 100644
--- a/src/frontend/tests/extended/features/lock-flow.spec.ts
+++ b/src/frontend/tests/extended/features/lock-flow.spec.ts
@@ -27,18 +27,21 @@ test(
       state: "visible",
     });
     await page.getByTestId("canvas_controls_dropdown").click();
+    await page.waitForTimeout(500);
 
     await page.getByTestId("flow_name").click();
     await page.getByTestId("lock-flow-switch").click();
-    await page.getByTestId("save-flow-settings").click();
-
+    await page.getByTestId("icon-Lock").isVisible({ timeout: 5000 });
     //ensure the UI is updated
     await page.waitForTimeout(500);
-
-    await page.waitForSelector('[data-testid="icon-Lock"]', {
-      timeout: 3000,
+    await page.getByTestId("save-flow-settings").click();
+    await expect(page.getByTestId("save-flow-settings")).toBeHidden({
+      timeout: 5000,
     });
 
+    //ensure the UI is updated
+    await page.getByTestId("icon-Lock").isVisible({ timeout: 5000 });
+
     await page.getByTestId("icon-ChevronLeft").click();
     await page.waitForSelector('[data-testid="mainpage_title"]', {
       timeout: 3000,
@@ -51,9 +54,9 @@ test(
       state: "visible",
     });
     await page.getByTestId("canvas_controls_dropdown").click();
+    await page.waitForTimeout(500);
 
     //ensure the UI is updated
-    await page.waitForTimeout(1000);
 
     await page.waitForSelector('[data-testid="icon-Lock"]', {
       timeout: 3000,
@@ -61,10 +64,15 @@ test(
 
     await page.getByTestId("flow_name").click();
     await page.getByTestId("lock-flow-switch").click();
-    await page.waitForSelector('[data-testid="icon-Lock"]', {
-      timeout: 3000,
-    });
+    await page.getByTestId("icon-Unlock").isVisible({ timeout: 5000 });
+    //ensure the UI is updated
+    await page.waitForTimeout(500);
     await page.getByTestId("save-flow-settings").click();
+    await expect(page.getByTestId("save-flow-settings")).toBeHidden({
+      timeout: 5000,
+    });
+    //ensure the UI is updated
+    await page.getByTestId("icon-Lock").isHidden({ timeout: 5000 });
 
     await page.getByTestId("icon-ChevronLeft").click();
     await page.waitForSelector('[data-testid="mainpage_title"]', {
@@ -79,6 +87,7 @@ test(
       state: "visible",
     });
     await page.getByTestId("canvas_controls_dropdown").click();
+    await page.waitForTimeout(500);
 
     await tryDeleteEdge(page);
     await page.locator(".react-flow__edge-path").nth(0).click();
@@ -125,8 +134,14 @@ test(
 async function tryConnectNodes(page: Page) {
   await page.getByTestId("flow_name").click();
   await page.getByTestId("lock-flow-switch").click();
-  await page.getByTestId("icon-Unlock").isVisible();
+  await page.getByTestId("lock-flow-switch").isHidden({ timeout: 5000 });
+  await page.getByTestId("icon-Unlock").isVisible({ timeout: 5000 });
+  //ensure the UI is updated
+  await page.waitForTimeout(500);
   await page.getByTestId("save-flow-settings").click();
+  await expect(page.getByTestId("save-flow-settings")).toBeHidden({
+    timeout: 5000,
+  });
 
   const numberOfTries = 5;
   let numberOfEdges = await page.locator(".react-flow__edge-path").count();
@@ -157,33 +172,56 @@ async function tryConnectNodes(page: Page) {
   await page.getByTestId("flow_name").click();
   await page.getByTestId("lock-flow-switch").click();
 
-  await page.getByTestId("lock_Unlock").isVisible();
+  await page.getByTestId("icon-Unlock").isVisible({ timeout: 5000 });
+  //ensure the UI is updated
+  await page.waitForTimeout(500);
   await page.getByTestId("save-flow-settings").click();
+  await expect(page.getByTestId("save-flow-settings")).toBeHidden({
+    timeout: 5000,
+  });
+
+  //ensure the UI is updated
+  await page.getByTestId("icon-Lock").isHidden({ timeout: 5000 });
 }
 
 async function tryDeleteEdge(page: Page) {
   await page.getByTestId("flow_name").click();
   await page.getByTestId("lock-flow-switch").click();
-  await page.getByTestId("icon-Lock").isVisible();
+  await page.getByTestId("icon-Lock").isVisible({ timeout: 5000 });
+  await page.waitForTimeout(500);
+
   await page.getByTestId("save-flow-settings").click();
+  await expect(page.getByTestId("save-flow-settings")).toBeHidden({
+    timeout: 5000,
+  });
+
+  //ensure the UI is updated
+  await page.getByTestId("icon-Lock").isVisible({ timeout: 5000 });
 
   let numberOfEdges = await page.locator(".react-flow__edge-path").count();
   expect(numberOfEdges).toBe(3);
-  const numberOfTries = 50;
+  const numberOfTries = 5;
 
   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();
+    await expect(
+      page.locator(".react-flow__edge-path").nth(0).click({ timeout: 500 }),
+    ).rejects.toThrow();
+    await expect(
+      page.locator(".react-flow__edge-path").nth(1).click({ timeout: 500 }),
+    ).rejects.toThrow();
+    await expect(
+      page.locator(".react-flow__edge-path").nth(2).click({ timeout: 500 }),
+    ).rejects.toThrow();
     expect(numberOfEdges).toBe(3);
   }
   //unlock the flow
   await page.getByTestId("flow_name").click();
   await page.getByTestId("lock-flow-switch").click();
-  await page.getByTestId("lock_Unlock").isVisible();
+  await page.getByTestId("icon-Unlock").isVisible({ timeout: 5000 });
+  //ensure the UI is updated
+  await page.waitForTimeout(500);
   await page.getByTestId("save-flow-settings").click();
+  await expect(page.getByTestId("save-flow-settings")).toBeHidden({
+    timeout: 5000,
+  });
 }