feat(apollo-react): loading state for toolbox

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

@@ -10,7 +10,7 @@ import {
 } from '@uipath/apollo-react/canvas/xyflow/react';
 import { ApButton, ApMenu } from '@uipath/apollo-react/material/components';
 import type { IMenuItem } from '@uipath/apollo-react/material/components/ap-menu/ApMenu.types';
-import { useCallback, useMemo, useState } from 'react';
+import { useCallback, useEffect, useMemo, useState } from 'react';
 import { DefaultCanvasTranslations } from '../../types';
 import {
   createGroupModificationHandlers,
@@ -1240,6 +1240,30 @@ export const AddAndReplaceTasks: Story = {
   args: {},
 };
 
+// Simulate async children fetch (2s delay)
+const fetchChildren = (id: string): Promise<ListItem[]> =>
+  new Promise((resolve) => {
+    setTimeout(() => {
+      resolve([
+        { id: `${id}-sub-1`, name: `${id} - Subtask A`, data: { type: `${id}-a` } },
+        { id: `${id}-sub-2`, name: `${id} - Subtask B`, data: { type: `${id}-b` } },
+        { id: `${id}-sub-3`, name: `${id} - Subtask C`, data: { type: `${id}-c` } },
+      ]);
+    }, 2000);
+  });
+
+// Top-level items with async children — level 2 loading is handled by Toolbox internally
+const loadedTaskOptionsWithChildren: ListItem[] = [
+  { id: 'email', name: 'Email', data: { type: 'email' }, children: (id) => fetchChildren(id) },
+  {
+    id: 'approval',
+    name: 'Approval',
+    data: { type: 'approval' },
+    children: (id) => fetchChildren(id),
+  },
+  { id: 'script', name: 'Script', data: { type: 'script' }, children: (id) => fetchChildren(id) },
+];
+
 const AddTaskLoadingStory = () => {
   const StageNodeWrapper = useMemo(
     () =>
@@ -1261,26 +1285,29 @@ const AddTaskLoadingStory = () => {
         width: 304,
         data: {
           stageDetails: {
-            label: 'Empty Stage (click +)',
+            label: 'Top-level loading (click +)',
             tasks: [],
           },
-          addTaskLoading: false,
+          // Top-level loading: skeletons until API returns items
+          addTaskLoading: true,
+          taskOptions: [] as ListItem[],
         },
       },
       {
-        id: 'loading-stage-with-tasks',
+        id: 'loading-stage-children',
         type: 'stage',
         position: { x: 400, y: 96 },
         width: 304,
         data: {
           stageDetails: {
-            label: 'With Tasks (click +)',
+            label: 'Async children (click +)',
             tasks: [
               [{ id: 'task-1', label: 'Existing Task', icon: <VerificationIcon /> }],
-              [{ id: 'task-2', label: 'Another Task', icon: <DocumentIcon /> }],
             ],
           },
+          // Children loading: items already available, level 2 loads on click
           addTaskLoading: false,
+          taskOptions: loadedTaskOptionsWithChildren,
         },
       },
     ],
@@ -1290,39 +1317,40 @@ const AddTaskLoadingStory = () => {
   const [nodesState, setNodes, onNodesChange] = useNodesState(initialNodes);
   const [edges, setEdges, onEdgesChange] = useEdgesState([]);
 
-  const setNodeLoading = useCallback(
-    (nodeId: string, loading: boolean) => {
+  // Simulate top-level API loading — after 3 seconds, items become available
+  useEffect(() => {
+    const timeout = setTimeout(() => {
       setNodes((nds) =>
         nds.map((node) =>
-          node.id === nodeId ? { ...node, data: { ...node.data, addTaskLoading: loading } } : node
+          node.id === 'loading-stage-empty'
+            ? {
+                ...node,
+                data: {
+                  ...node.data,
+                  addTaskLoading: false,
+                  taskOptions: loadedTaskOptionsWithChildren,
+                },
+              }
+            : node
         )
       );
-    },
-    [setNodes]
-  );
+    }, 3000);
+    return () => clearTimeout(timeout);
+  }, [setNodes]);
 
-  const handleTaskAddForNode = useCallback(
-    (nodeId: string) => {
-      setNodeLoading(nodeId, true);
-      // Simulate API delay of 3 seconds
-      setTimeout(() => {
-        setNodeLoading(nodeId, false);
-      }, 3000);
-    },
-    [setNodeLoading]
-  );
-
-  // Inject a per-node onTaskAdd handler
+  // Inject per-node handlers
   const nodesWithHandler = useMemo(
     () =>
       nodesState.map((node) => ({
         ...node,
         data: {
           ...node.data,
-          onTaskAdd: () => handleTaskAddForNode(node.id),
+          onAddTaskFromToolbox: (taskItem: ListItem) => {
+            window.alert(`Added task "${taskItem.name}" to ${node.id}`);
+          },
         },
       })),
-    [nodesState, handleTaskAddForNode]
+    [nodesState]
   );
 
   const onConnect = useCallback(

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

@@ -71,15 +71,17 @@ vi.mock('../Toolbox', () => ({
   Toolbox: ({
     title,
     initialItems,
+    loading,
     onItemSelect,
     onClose,
   }: {
     title: string;
     initialItems: ListItem[];
+    loading?: boolean;
     onItemSelect?: (item: ListItem) => void;
     onClose?: () => void;
   }) => (
-    <div data-testid="toolbox">
+    <div data-testid="toolbox" data-loading={loading}>
       <div data-testid="toolbox-title">{title}</div>
       <div data-testid="toolbox-items-count">{initialItems.length}</div>
       {initialItems.map((item, index) => (
@@ -614,63 +616,54 @@ describe('StageNode - Add Task Loading State', () => {
     vi.clearAllMocks();
   });
 
-  it('should show the add icon by default when addTaskLoading is not set', () => {
-    const onTaskAdd = vi.fn();
-    renderStageNode({ onTaskAdd });
+  it('should pass addTaskLoading to Toolbox when toolbox is open', async () => {
+    const user = userEvent.setup();
+    const onAddTaskFromToolbox = vi.fn();
+    renderStageNode({ onAddTaskFromToolbox, addTaskLoading: true });
 
-    // The add icon should be present, no spinner
-    expect(screen.queryByTestId('ap-circular-progress')).not.toBeInTheDocument();
-  });
-
-  it('should show a loading spinner instead of the add icon when addTaskLoading is true', () => {
-    const onTaskAdd = vi.fn();
-    renderStageNode({ onTaskAdd, addTaskLoading: true });
-
-    expect(screen.getByTestId('ap-circular-progress')).toBeInTheDocument();
-  });
-
-  it('should disable the add task button when addTaskLoading is true', () => {
-    const onTaskAdd = vi.fn();
-    renderStageNode({ onTaskAdd, addTaskLoading: true });
+    const addButton = screen.getByRole('button', { name: 'Add task' });
+    await user.click(addButton);
 
-    const spinner = screen.getByTestId('ap-circular-progress');
-    // The disabled button is the closest button ancestor of the spinner
-    const button = spinner.closest('button');
-    expect(button).toBeDisabled();
+    const toolbox = screen.getByTestId('toolbox');
+    expect(toolbox).toHaveAttribute('data-loading', 'true');
   });
 
-  it('should not call onTaskAdd when button is disabled while loading', () => {
-    const onTaskAdd = vi.fn();
-    renderStageNode({ onTaskAdd, addTaskLoading: true });
+  it('should not pass loading to Toolbox when addTaskLoading is false', async () => {
+    const user = userEvent.setup();
+    const onAddTaskFromToolbox = vi.fn();
+    renderStageNode({ onAddTaskFromToolbox, addTaskLoading: false });
 
-    const spinner = screen.getByTestId('ap-circular-progress');
-    const button = spinner.closest('button') as HTMLButtonElement;
+    const addButton = screen.getByRole('button', { name: 'Add task' });
+    await user.click(addButton);
 
-    // Button is disabled and has pointer-events: none, preventing any clicks
-    expect(button).toBeDisabled();
-    button.click();
-    expect(onTaskAdd).not.toHaveBeenCalled();
+    const toolbox = screen.getByTestId('toolbox');
+    expect(toolbox).toHaveAttribute('data-loading', 'false');
   });
 
-  it('should show the add icon when addTaskLoading is false', () => {
-    const onTaskAdd = vi.fn();
-    renderStageNode({ onTaskAdd, addTaskLoading: false });
+  it('should not disable the add button when addTaskLoading is true', () => {
+    const onAddTaskFromToolbox = vi.fn();
+    renderStageNode({ onAddTaskFromToolbox, addTaskLoading: true });
 
-    expect(screen.queryByTestId('ap-circular-progress')).not.toBeInTheDocument();
+    const addButton = screen.getByRole('button', { name: 'Add task' });
+    expect(addButton).not.toBeDisabled();
   });
 
-  it('should switch from spinner back to add icon when addTaskLoading changes to false', () => {
-    const onTaskAdd = vi.fn();
-    const { rerender } = renderStageNode({ onTaskAdd, addTaskLoading: true });
+  it('should update Toolbox loading when addTaskLoading changes to false', async () => {
+    const user = userEvent.setup();
+    const onAddTaskFromToolbox = vi.fn();
+    const { rerender } = renderStageNode({ onAddTaskFromToolbox, addTaskLoading: true });
+
+    const addButton = screen.getByRole('button', { name: 'Add task' });
+    await user.click(addButton);
 
-    expect(screen.getByTestId('ap-circular-progress')).toBeInTheDocument();
+    expect(screen.getByTestId('toolbox')).toHaveAttribute('data-loading', 'true');
 
     rerender(
       <ReactFlowProvider>
-        <StageNode {...defaultProps} onTaskAdd={onTaskAdd} addTaskLoading={false} />
+        <StageNode {...defaultProps} onAddTaskFromToolbox={onAddTaskFromToolbox} addTaskLoading={false} />
       </ReactFlowProvider>
     );
 
-    expect(screen.queryByTestId('ap-circular-progress')).not.toBeInTheDocument();
+    expect(screen.getByTestId('toolbox')).toHaveAttribute('data-loading', 'false');
   });
 });

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

@@ -20,7 +20,6 @@ import { FontVariantToken, Padding, Spacing } from '@uipath/apollo-core';
 import { Column, Row } from '@uipath/apollo-react/canvas/layouts';
 import { Position, useStore, useViewport } from '@uipath/apollo-react/canvas/xyflow/react';
 import {
-  ApCircularProgress,
   ApIcon,
   ApIconButton,
   ApLink,
@@ -543,19 +542,14 @@ const StageNodeComponent = (props: StageNodeProps) => {
               </ApTooltip>
             )}
             {(onTaskAdd || onAddTaskFromToolbox) && !isReadOnly && (
-              <ApTooltip content={addTaskLoading ? 'Loading...' : addTaskLabel} placement="top">
+              <ApTooltip content={addTaskLabel} placement="top">
                 <span>
                   <ApIconButton
                     onClick={handleTaskAddClick}
                     size="small"
-                    disabled={addTaskLoading}
                     label={addTaskLabel}
                   >
-                    {addTaskLoading ? (
-                      <ApCircularProgress size={20} />
-                    ) : (
-                      <ApIcon name="add" size="20px" />
-                    )}
+                    <ApIcon name="add" size="20px" />
                   </ApIconButton>
                 </span>
               </ApTooltip>
@@ -682,6 +676,7 @@ const StageNodeComponent = (props: StageNodeProps) => {
           <Toolbox
             title={addTaskLabel}
             initialItems={taskOptions}
+            loading={addTaskLoading}
             onClose={() => setIsAddingTask(false)}
             onItemSelect={handleAddTaskToolboxItemSelected}
             onSearch={onTaskToolboxSearch}