refactor(apollo-react): task replace in case mgmt [MST-7038]

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

@@ -966,9 +966,7 @@ const AddAndReplaceTasksStory = () => {
   const nodeTypes = useMemo(() => ({ stage: StageNodeWrapper }), [StageNodeWrapper]);
   const edgeTypes = useMemo(() => ({ stage: StageEdge }), []);
 
-  const [pendingReplaceTask, setPendingReplaceTask] = useState<
-    { groupIndex: number; taskIndex: number } | undefined
-  >();
+  const [pendingReplaceTask, setPendingReplaceTask] = useState(false);
   const [selectedTaskId, setSelectedTaskId] = useState<string | undefined>();
   const [menuAnchorEl, setMenuAnchorEl] = useState<null | HTMLElement>(null);
 
@@ -1071,7 +1069,8 @@ const AddAndReplaceTasksStory = () => {
         })
       );
 
-      setPendingReplaceTask(undefined);
+      setPendingReplaceTask(false);
+      setSelectedTaskId(undefined);
     },
     [setNodes]
   );
@@ -1118,20 +1117,41 @@ const AddAndReplaceTasksStory = () => {
     setSelectedTaskId(taskId);
   }, []);
 
+  // Clear task selection when clicking on the canvas background
+  const handlePaneClick = useCallback(() => {
+    setSelectedTaskId(undefined);
+    setPendingReplaceTask(false);
+  }, []);
+
+  // Clear replace state when the node is deselected (e.g., clicking on canvas)
+  const handleNodesChange = useCallback(
+    (...args: Parameters<typeof onNodesChange>) => {
+      onNodesChange(...args);
+      const deselected = args[0].some(
+        (c) => c.type === 'select' && c.id === 'add-replace-stage' && !c.selected
+      );
+      if (deselected) {
+        setPendingReplaceTask(false);
+        setSelectedTaskId(undefined);
+      }
+    },
+    [onNodesChange]
+  );
+
   // Get current tasks from node state for menu items
   const currentTasks =
     nodesState.find((n) => n.id === 'add-replace-stage')?.data.stageDetails.tasks || [];
 
   // Create menu items for task selection
   const taskMenuItems = useMemo<IMenuItem[]>(() => {
-    return currentTasks.flatMap((group, groupIndex) =>
-      group.map((task, taskIndex) => ({
+    return currentTasks.flatMap((group) =>
+      group.map((task) => ({
         title: task.label,
         variant: 'item' as const,
         startIcon: task.icon,
         onClick: () => {
           setSelectedTaskId(task.id);
-          setPendingReplaceTask({ groupIndex, taskIndex });
+          setPendingReplaceTask(true);
           setMenuAnchorEl(null);
         },
       }))
@@ -1147,10 +1167,10 @@ const AddAndReplaceTasksStory = () => {
               ...node,
               data: {
                 ...node.data,
-                pendingReplaceTask: pendingReplaceTask,
+                ...(selectedTaskId && { pendingReplaceTask }),
                 stageDetails: {
                   ...node.data.stageDetails,
-                  selectedTasks: selectedTaskId ? [selectedTaskId] : undefined,
+                  selectedTaskId,
                 },
                 onAddTaskFromToolbox: handleAddTask,
                 onReplaceTaskFromToolbox: handleReplaceTask,
@@ -1178,25 +1198,25 @@ const AddAndReplaceTasksStory = () => {
 
   // Compute button label based on whether a task is being replaced
   const replaceButtonLabel = useMemo(() => {
-    if (pendingReplaceTask) {
-      const taskBeingReplaced =
-        currentTasks[pendingReplaceTask.groupIndex]?.[pendingReplaceTask.taskIndex];
+    if (pendingReplaceTask && selectedTaskId) {
+      const taskBeingReplaced = currentTasks.flat().find((t) => t.id === selectedTaskId);
       if (taskBeingReplaced) {
         return `Replacing Task: ${taskBeingReplaced.label}`;
       }
     }
     return 'Replace Task';
-  }, [pendingReplaceTask, currentTasks]);
+  }, [pendingReplaceTask, selectedTaskId, currentTasks]);
 
   return (
     <div style={{ width: '100vw', height: '100vh' }}>
       <ReactFlowProvider>
         <BaseCanvas
           nodes={nodesWithMetadata}
           edges={edges}
-          onNodesChange={onNodesChange}
+          onNodesChange={handleNodesChange}
           onEdgesChange={onEdgesChange}
           onConnect={onConnect}
+          onPaneClick={handlePaneClick}
           nodeTypes={nodeTypes}
           edgeTypes={edgeTypes}
           mode="design"

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

@@ -96,7 +96,7 @@ const StageNodeComponent = (props: StageNodeProps) => {
   const isException = stageDetails?.isException;
   const isReadOnly = !!stageDetails?.isReadOnly;
   const icon = stageDetails?.icon;
-  const selectedTasks = stageDetails?.selectedTasks;
+  const selectedTaskId = stageDetails?.selectedTaskId;
   const defaultContent = stageDetails?.defaultContent || 'Add first task';
 
   const status = execution?.stageStatus?.status;
@@ -130,16 +130,21 @@ const StageNodeComponent = (props: StageNodeProps) => {
   const [isReplacingTask, setIsReplacingTask] = useState(false);
 
   useEffect(() => {
-    if (pendingReplaceTask?.groupIndex != null && pendingReplaceTask?.taskIndex != null) {
-      const taskGroup = tasks[pendingReplaceTask.groupIndex];
-      taskStateReference.current = {
-        isParallel: (taskGroup?.length ?? 0) > 1,
-        groupIndex: pendingReplaceTask.groupIndex,
-        taskIndex: pendingReplaceTask.taskIndex,
-      };
-      setIsReplacingTask(true);
-    } else setIsReplacingTask(false);
-  }, [pendingReplaceTask, tasks]);
+    if (pendingReplaceTask) {
+      const match = tasks
+        .flatMap((group, gi) => group.map((task, ti) => ({ task, groupIndex: gi, taskIndex: ti })))
+        .find(({ task }) => task.id === selectedTaskId);
+
+      if (match) {
+        taskStateReference.current = {
+          isParallel: (tasks[match.groupIndex]?.length ?? 0) > 1,
+          groupIndex: match.groupIndex,
+          taskIndex: match.taskIndex,
+        };
+        setIsReplacingTask(true);
+      }
+    }
+  }, [pendingReplaceTask, selectedTaskId, tasks]);
 
   const [activeDragId, setActiveDragId] = useState<string | null>(null);
   const [offsetLeft, setOffsetLeft] = useState(0);
@@ -245,7 +250,13 @@ const StageNodeComponent = (props: StageNodeProps) => {
       const items: NodeMenuItem[] = [];
 
       if (onReplaceTaskFromToolbox) {
-        items.push(getMenuItem('replace-task', 'Replace task', () => setIsReplacingTask(true)));
+        items.push(
+          getMenuItem('replace-task', 'Replace task', () => {
+            const taskId = tasks[groupIndex]?.[taskIndex]?.id;
+            if (taskId) onTaskClick?.(taskId);
+            setIsReplacingTask(true);
+          })
+        );
         items.push(getDivider());
       }
 
@@ -265,7 +276,7 @@ const StageNodeComponent = (props: StageNodeProps) => {
 
       return items;
     },
-    [onReplaceTaskFromToolbox, onTaskGroupModification, reGroupTaskFunction]
+    [onReplaceTaskFromToolbox, onTaskClick, onTaskGroupModification, reGroupTaskFunction, tasks]
   );
 
   const { setSelectedNodeId } = useNodeSelection();
@@ -306,21 +317,14 @@ const StageNodeComponent = (props: StageNodeProps) => {
 
   const handleReplaceTaskToolboxItemSelected = useCallback(
     (item: ListItem) => {
-      setIsReplacingTask(true);
-      const groupIndex = taskStateReference.current.groupIndex;
-      const taskIndex = taskStateReference.current.taskIndex;
-      const taskId = tasks[groupIndex]?.[taskIndex]?.id;
-      if (taskId) {
-        onTaskClick?.(taskId);
-      }
       onReplaceTaskFromToolbox?.(
         item,
         taskStateReference.current.groupIndex,
         taskStateReference.current.taskIndex
       );
       setIsReplacingTask(false);
     },
-    [onReplaceTaskFromToolbox, onTaskClick, tasks]
+    [onReplaceTaskFromToolbox]
   );
 
   const handleConfigurations: HandleGroupManifest[] = useMemo(
@@ -619,7 +623,7 @@ const StageNodeComponent = (props: StageNodeProps) => {
                                 key={task.id}
                                 task={task}
                                 taskExecution={taskExecution}
-                                isSelected={!!selectedTasks?.includes(task.id)}
+                                isSelected={selectedTaskId === task.id}
                                 isParallel={isParallel}
                                 contextMenuItems={contextMenuItems(
                                   isParallel,

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

@@ -27,10 +27,7 @@ export interface StageNodeProps extends NodeProps {
   dragging: boolean;
   selected: boolean;
   id: string;
-  pendingReplaceTask?: {
-    groupIndex: number;
-    taskIndex: number;
-  };
+  pendingReplaceTask?: boolean;
   stageDetails: {
     label: string;
     defaultContent?: string;
@@ -42,7 +39,7 @@ export interface StageNodeProps extends NodeProps {
     isException?: boolean;
     isReadOnly?: boolean;
     tasks: StageTaskItem[][];
-    selectedTasks?: string[];
+    selectedTaskId?: string;
   };
   addTaskLabel?: string;
   addTaskLoading?: boolean;

packages/apollo-react/src/canvas/components/Toolbox/Toolbox.tsx

@@ -101,6 +101,7 @@ export function Toolbox<T>({
     parentItem: ListItem<T> | null;
   }>();
 
+  const containerRef = useRef<HTMLDivElement>(null);
   const transitionTimeoutRef = useRef<ReturnType<typeof setTimeout> | undefined>(undefined);
   const searchIdRef = useRef(0);
   const initialItemsRef = useRef(initialItems);
@@ -297,8 +298,18 @@ export function Toolbox<T>({
       }
     };
 
+    const handleClickOutside = (e: MouseEvent) => {
+      if (containerRef.current && !containerRef.current.contains(e.target as Node)) {
+        onClose();
+      }
+    };
+
     document.addEventListener('keydown', handleKeyDown);
-    return () => document.removeEventListener('keydown', handleKeyDown);
+    document.addEventListener('mousedown', handleClickOutside);
+    return () => {
+      document.removeEventListener('keydown', handleKeyDown);
+      document.removeEventListener('mousedown', handleClickOutside);
+    };
   }, [
     isSearching,
     navigationStack.canGoBack,
@@ -309,27 +320,29 @@ export function Toolbox<T>({
   ]);
 
   return (
-    <Column px={20} py={12} gap={12} w={320} h={440}>
-      <Header
-        title={currentParentItem?.name || title}
-        onBack={handleBackTransition}
-        showBackButton={navigationStack.canGoBack}
-      />
-
-      <SearchBox value={search} onChange={handleSearch} clear={clearSearch} placeholder="Search" />
-
-      <AnimatedContainer>
-        <AnimatedContent entering={isTransitioning} direction={animationDirection}>
-          <ListView
-            isLoading={childrenLoading || searchLoading || loading}
-            items={isSearching && !isSearchingInitialItems ? searchedItems : items}
-            emptyStateMessage={isSearching ? 'No matching nodes found' : 'No nodes found'}
-            enableSections={!isSearching}
-            onItemClick={handleItemSelect}
-            onItemHover={onItemHover}
-          />
-        </AnimatedContent>
-      </AnimatedContainer>
-    </Column>
+    <div ref={containerRef}>
+      <Column px={20} py={12} gap={12} w={320} h={440}>
+        <Header
+          title={currentParentItem?.name || title}
+          onBack={handleBackTransition}
+          showBackButton={navigationStack.canGoBack}
+        />
+
+        <SearchBox value={search} onChange={handleSearch} clear={clearSearch} placeholder="Search" />
+
+        <AnimatedContainer>
+          <AnimatedContent entering={isTransitioning} direction={animationDirection}>
+            <ListView
+              isLoading={childrenLoading || searchLoading || loading}
+              items={isSearching && !isSearchingInitialItems ? searchedItems : items}
+              emptyStateMessage={isSearching ? 'No matching nodes found' : 'No nodes found'}
+              enableSections={!isSearching}
+              onItemClick={handleItemSelect}
+              onItemHover={onItemHover}
+            />
+          </AnimatedContent>
+        </AnimatedContainer>
+      </Column>
+    </div>
   );
 }