fix(apollo-react): disable task three-dot menu while loading [MST-8171]

packages/apollo-react/src/canvas/components/StageNode/AdhocTask.test.tsx

@@ -293,4 +293,76 @@ describe('AdhocTaskItem', () => {
       });
     });
   });
+
+  describe('Task Loading State', () => {
+    it('disables menu button when isTaskLoading is true', () => {
+      const onRemove = vi.fn();
+      const menuItems = createMenuItems(onRemove);
+
+      render(
+        <AdhocTaskItem
+          {...defaultProps}
+          getContextMenuItems={() => menuItems}
+          isTaskLoading={true}
+        />
+      );
+
+      const menuButton = screen.getByTestId('stage-task-menu-adhoc-1');
+      expect(menuButton).toBeDisabled();
+    });
+
+    it('does not disable menu button when isTaskLoading is false', () => {
+      const onRemove = vi.fn();
+      const menuItems = createMenuItems(onRemove);
+
+      render(
+        <AdhocTaskItem
+          {...defaultProps}
+          getContextMenuItems={() => menuItems}
+          isTaskLoading={false}
+        />
+      );
+
+      const menuButton = screen.getByTestId('stage-task-menu-adhoc-1');
+      expect(menuButton).not.toBeDisabled();
+    });
+
+    it('does not open menu when clicking a disabled menu button', async () => {
+      const user = userEvent.setup();
+      const onRemove = vi.fn();
+      const menuItems = createMenuItems(onRemove);
+
+      render(
+        <AdhocTaskItem
+          {...defaultProps}
+          getContextMenuItems={() => menuItems}
+          isTaskLoading={true}
+        />
+      );
+
+      const menuButton = screen.getByTestId('stage-task-menu-adhoc-1');
+      await user.click(menuButton);
+
+      expect(screen.queryByText('Replace task')).not.toBeInTheDocument();
+    });
+
+    it('does not open menu on right-click when isTaskLoading is true', async () => {
+      const user = userEvent.setup();
+      const onRemove = vi.fn();
+      const menuItems = createMenuItems(onRemove);
+
+      render(
+        <AdhocTaskItem
+          {...defaultProps}
+          getContextMenuItems={() => menuItems}
+          isTaskLoading={true}
+        />
+      );
+
+      const task = screen.getByTestId('stage-task-adhoc-1');
+      await user.pointer({ keys: '[MouseRight]', target: task });
+
+      expect(screen.queryByText('Replace task')).not.toBeInTheDocument();
+    });
+  });
 });

packages/apollo-react/src/canvas/components/StageNode/AdhocTask.tsx

@@ -71,6 +71,7 @@ interface AdhocTaskItemProps {
   getContextMenuItems?: () => NodeMenuItem[];
   onTaskClick: (e: React.MouseEvent, taskId: string) => void;
   onTaskPlay?: (taskId: string) => Promise<void>;
+  isTaskLoading?: boolean;
 }
 
 const AdhocTaskItemComponent = ({
@@ -80,6 +81,7 @@ const AdhocTaskItemComponent = ({
   getContextMenuItems,
   onTaskClick,
   onTaskPlay,
+  isTaskLoading,
 }: AdhocTaskItemProps) => {
   const [isMenuOpen, setIsMenuOpen] = useState(false);
   const taskRef = useRef<HTMLDivElement>(null);
@@ -108,7 +110,7 @@ const AdhocTaskItemComponent = ({
       selected={isSelected}
       status={taskExecution?.status}
       onClick={handleClick}
-      {...(getContextMenuItems && { onContextMenu: handleContextMenu })}
+      {...(getContextMenuItems && !isTaskLoading && { onContextMenu: handleContextMenu })}
     >
       <TaskContent task={task} taskExecution={taskExecution} />
       {onTaskPlay && <AdhocTaskPlayButton taskId={task.id} onTaskPlay={onTaskPlay} />}
@@ -119,6 +121,7 @@ const AdhocTaskItemComponent = ({
           getContextMenuItems={getContextMenuItems}
           onMenuOpenChange={handleMenuOpenChange}
           taskRef={taskRef}
+          disabled={isTaskLoading}
         />
       )}
     </StageTask>

packages/apollo-react/src/canvas/components/StageNode/DraggableTask.test.tsx

@@ -346,6 +346,78 @@ describe('DraggableTask', () => {
     });
   });
 
+  describe('Task Loading State', () => {
+    it('disables menu button when isTaskLoading is true', () => {
+      const onRemove = vi.fn();
+      const menuItems = createMenuItems(onRemove);
+
+      render(
+        <DraggableTask
+          {...defaultProps}
+          getContextMenuItems={() => menuItems}
+          isTaskLoading={true}
+        />
+      );
+
+      const menuButton = screen.getByTestId('stage-task-menu-task-1');
+      expect(menuButton).toBeDisabled();
+    });
+
+    it('does not disable menu button when isTaskLoading is false', () => {
+      const onRemove = vi.fn();
+      const menuItems = createMenuItems(onRemove);
+
+      render(
+        <DraggableTask
+          {...defaultProps}
+          getContextMenuItems={() => menuItems}
+          isTaskLoading={false}
+        />
+      );
+
+      const menuButton = screen.getByTestId('stage-task-menu-task-1');
+      expect(menuButton).not.toBeDisabled();
+    });
+
+    it('does not open menu when clicking a disabled menu button', async () => {
+      const user = userEvent.setup();
+      const onRemove = vi.fn();
+      const menuItems = createMenuItems(onRemove);
+
+      render(
+        <DraggableTask
+          {...defaultProps}
+          getContextMenuItems={() => menuItems}
+          isTaskLoading={true}
+        />
+      );
+
+      const menuButton = screen.getByTestId('stage-task-menu-task-1');
+      await user.click(menuButton);
+
+      expect(screen.queryByText('Move Up')).not.toBeInTheDocument();
+    });
+
+    it('does not open menu on right-click when isTaskLoading is true', async () => {
+      const user = userEvent.setup();
+      const onRemove = vi.fn();
+      const menuItems = createMenuItems(onRemove);
+
+      render(
+        <DraggableTask
+          {...defaultProps}
+          getContextMenuItems={() => menuItems}
+          isTaskLoading={true}
+        />
+      );
+
+      const task = screen.getByTestId('stage-task-task-1');
+      await user.pointer({ keys: '[MouseRight]', target: task });
+
+      expect(screen.queryByText('Move Up')).not.toBeInTheDocument();
+    });
+  });
+
   describe('Task Rendering', () => {
     it('renders task label', () => {
       render(<DraggableTask {...defaultProps} />);

packages/apollo-react/src/canvas/components/StageNode/DraggableTask.tsx

@@ -131,6 +131,7 @@ const DraggableTaskComponent = ({
   onTaskClick,
   isDragDisabled,
   projectedDepth,
+  isTaskLoading,
 }: DraggableTaskProps) => {
   const [isMenuOpen, setIsMenuOpen] = useState(false);
   const zoom = useStore((s) => s.transform[2]);
@@ -208,7 +209,7 @@ const DraggableTaskComponent = ({
       isParallel={isParallel}
       isDragEnabled={!isDragDisabled}
       onClick={handleClick}
-      {...(getContextMenuItems && { onContextMenu: handleContextMenu })}
+      {...(getContextMenuItems && !isTaskLoading && { onContextMenu: handleContextMenu })}
     >
       <TaskContent task={task} taskExecution={taskExecution} />
 
@@ -233,6 +234,7 @@ const DraggableTaskComponent = ({
           getContextMenuItems={handleGetContextMenuItems}
           onMenuOpenChange={handleMenuOpenChange}
           taskRef={taskRef}
+          disabled={isTaskLoading}
         />
       )}
     </StageTask>

packages/apollo-react/src/canvas/components/StageNode/DraggableTask.types.ts

@@ -19,4 +19,5 @@ export interface DraggableTaskProps {
   onTaskPlay?: (taskId: string) => Promise<void>;
   isDragDisabled?: boolean;
   projectedDepth?: number;
+  isTaskLoading?: boolean;
 }

packages/apollo-react/src/canvas/components/StageNode/StageNode.stories.tsx

@@ -1443,18 +1443,57 @@ const AddTaskLoadingStory = () => {
   const setNodesRef = useRef<React.Dispatch<React.SetStateAction<Node[]>>>(null!);
 
   // Inject per-node handlers — simulates loading state on add-task:
-  // 1. Set addTaskLoading=true so the + button is disabled
-  // 2. After 2s, set addTaskLoading=false to re-enable it
+  // 1. Set loadingTaskIds with a placeholder ID so the + button is disabled
+  // 2. After 2s, clear loadingTaskIds to re-enable it
   const timeoutRef = useRef<ReturnType<typeof setTimeout> | undefined>(undefined);
 
-  const handleAddTaskFromToolbox = useCallback((nodeId: string, _taskItem: ListItem) => {
+  const handleAddTaskFromToolbox = useCallback((nodeId: string, taskItem: ListItem) => {
     clearTimeout(timeoutRef.current);
+    const newTaskId = `task-${Date.now()}`;
+    // Add the task to the stage and mark it as loading
     setNodesRef.current((nds) =>
-      nds.map((n) => (n.id === nodeId ? { ...n, data: { ...n.data, addTaskLoading: true } } : n))
+      nds.map((n) => {
+        if (n.id !== nodeId) return n;
+        const data = n.data as Record<string, any>;
+        const currentTasks: StageTaskItem[][] = data.stageDetails?.tasks ?? [];
+        const newTask: StageTaskItem = { id: newTaskId, label: taskItem.name };
+        const currentExecution = data.execution ?? {
+          stageStatus: { status: undefined },
+          taskStatus: {},
+        };
+        return {
+          ...n,
+          data: {
+            ...data,
+            stageDetails: { ...data.stageDetails, tasks: [...currentTasks, [newTask]] },
+            loadingTaskIds: new Set([...((data.loadingTaskIds as Set<string>) ?? []), newTaskId]),
+            execution: {
+              ...currentExecution,
+              taskStatus: { ...currentExecution.taskStatus, [newTaskId]: { status: 'InProgress' } },
+            },
+          },
+        };
+      })
     );
+    // After 2s, clear loading state and execution status
     timeoutRef.current = setTimeout(() => {
       setNodesRef.current((nds) =>
-        nds.map((n) => (n.id === nodeId ? { ...n, data: { ...n.data, addTaskLoading: false } } : n))
+        nds.map((n) => {
+          if (n.id !== nodeId) return n;
+          const data = n.data as Record<string, any>;
+          const { [newTaskId]: _, ...remainingTaskStatus } = data.execution?.taskStatus ?? {};
+          return {
+            ...n,
+            data: {
+              ...data,
+              loadingTaskIds: new Set(),
+              execution: {
+                ...data.execution,
+                taskStatus: remainingTaskStatus,
+              },
+            },
+          };
+        })
       );
     }, 2000);
   }, []);
@@ -1471,11 +1510,12 @@ const AddTaskLoadingStory = () => {
             label: 'Loading (+ disabled for 3s)',
             tasks: [],
           },
-          addTaskLoading: true,
+          loadingTaskIds: new Set(['loading-task']),
           taskOptions: [] as ListItem[],
           onAddTaskFromToolbox: (taskItem: ListItem) => {
             handleAddTaskFromToolbox('loading-stage-empty', taskItem);
           },
+          onTaskGroupModification: () => {},
         },
       },
       {
@@ -1488,11 +1528,45 @@ const AddTaskLoadingStory = () => {
             label: 'Async children (click +)',
             tasks: [[{ id: 'task-1', label: 'Existing Task', icon: <VerificationIcon /> }]],
           },
-          addTaskLoading: false,
+          loadingTaskIds: new Set(),
           taskOptions: loadedTaskOptionsWithChildren,
           onAddTaskFromToolbox: (taskItem: ListItem) => {
             handleAddTaskFromToolbox('loading-stage-children', taskItem);
           },
+          onTaskGroupModification: () => {},
+        },
+      },
+      {
+        id: 'loading-stage-tasks',
+        type: 'stage',
+        position: { x: 752, y: 96 },
+        width: 304,
+        data: {
+          stageDetails: {
+            label: 'Task loading (3-dot disabled)',
+            tasks: [
+              [{ id: 'loading-task-1', label: 'Loading Task (3-dot disabled)' }],
+              [
+                {
+                  id: 'ready-task-1',
+                  label: 'Ready Task (3-dot enabled)',
+                  icon: <VerificationIcon />,
+                },
+              ],
+            ],
+          },
+          loadingTaskIds: new Set(['loading-task-1']),
+          execution: {
+            stageStatus: { status: undefined },
+            taskStatus: {
+              'loading-task-1': { status: 'InProgress' },
+            },
+          },
+          taskOptions: loadedTaskOptionsWithChildren,
+          onAddTaskFromToolbox: (taskItem: ListItem) => {
+            handleAddTaskFromToolbox('loading-stage-tasks', taskItem);
+          },
+          onTaskGroupModification: () => {},
         },
       },
     ],
@@ -1513,7 +1587,7 @@ const AddTaskLoadingStory = () => {
                 ...node,
                 data: {
                   ...node.data,
-                  addTaskLoading: false,
+                  loadingTaskIds: new Set(),
                   taskOptions: loadedTaskOptionsWithChildren,
                 },
               }
@@ -1524,6 +1598,30 @@ const AddTaskLoadingStory = () => {
     return () => clearTimeout(timeout);
   }, [setNodes]);
 
+  // Simulate per-task loading — after 5 seconds, task finishes loading and 3-dot becomes enabled
+  useEffect(() => {
+    const timeout = setTimeout(() => {
+      setNodes((nds) =>
+        nds.map((node) =>
+          node.id === 'loading-stage-tasks'
+            ? {
+                ...node,
+                data: {
+                  ...node.data,
+                  loadingTaskIds: new Set(),
+                  execution: {
+                    stageStatus: { status: undefined },
+                    taskStatus: {},
+                  },
+                },
+              }
+            : node
+        )
+      );
+    }, 5000);
+    return () => clearTimeout(timeout);
+  }, [setNodes]);
+
   useEffect(() => {
     return () => clearTimeout(timeoutRef.current);
   }, []);