feature(FR-352): update-useBAINotification

react/src/components/ComputeSessionNodeItems/ContainerCommitModal.tsx

@@ -58,7 +58,7 @@ const ContainerCommitModal: React.FC<ContainerCommitModalProps> = ({
               backgroundTask: {
                 taskId: res.task_id,
                 status: 'pending',
-                statusDescriptions: {
+                onChange: {
                   pending: t('session.CommitOnGoing'),
                   resolved: t('session.CommitFinished'),
                   rejected: t('session.CommitFailed'),

react/src/components/ContainerRegistryList.tsx

@@ -221,7 +221,7 @@ const ContainerRegistryList: React.FC<{
               status: 'pending',
               percent: 0,
               taskId: rescan_images.task_id,
-              statusDescriptions: {
+              onChange: {
                 pending: t('registry.RescanImages'),
                 resolved: t('registry.RegistryUpdateFinished'),
                 rejected: t('registry.RegistryUpdateFailed'),

react/src/components/ContainerRegistryListBefore2409.tsx

@@ -176,7 +176,7 @@ const ContainerRegistryListBefore2409: React.FC<{
               status: 'pending',
               percent: 0,
               taskId: rescan_images.task_id,
-              statusDescriptions: {
+              onChange: {
                 pending: t('registry.RescanImages'),
                 resolved: t('registry.RegistryUpdateFinished'),
                 rejected: t('registry.RegistryUpdateFailed'),

react/src/components/MaintenanceSettingList.tsx

@@ -1,8 +1,5 @@
 import { useSuspendedBackendaiClient } from '../hooks';
-import {
-  CLOSING_DURATION,
-  useSetBAINotification,
-} from '../hooks/useBAINotification';
+import { useSetBAINotification } from '../hooks/useBAINotification';
 import SettingList, { SettingGroup } from './SettingList';
 import { RedoOutlined } from '@ant-design/icons';
 import { Button } from 'antd';
@@ -29,7 +26,7 @@ const MaintenanceSettingList = () => {
         promise: baiClient.maintenance
           .recalculate_usage()
           .finally(() => setIsRecalculating(false)),
-        statusDescriptions: {
+        onChange: {
           pending: t('maintenance.Recalculating'),
           resolved: t('maintenance.RecalculationFinished'),
           rejected: t('maintenance.RecalculationFailed'),
@@ -40,54 +37,45 @@ const MaintenanceSettingList = () => {
 
   const rescanImages = () => {
     setIsRescanning(true);
-    const notiKey = upsertNotification({
+    upsertNotification({
       message: t('maintenance.RescanImages'),
-      description: t('maintenance.RescanImageScanning'),
       open: true,
       backgroundTask: {
         status: 'pending',
+        promise: baiClient.maintenance.rescan_images(),
+        onChange: {
+          pending: t('maintenance.RescanImageScanning'),
+          resolved: (data) => {
+            const result = data as Awaited<
+              ReturnType<typeof baiClient.maintenance.rescan_images>
+            >;
+            if (result.rescan_images.ok) {
+              return {
+                backgroundTask: {
+                  status: 'pending',
+                  taskId: result.rescan_images.task_id,
+                  promise: null,
+                  percent: 0,
+                  onChange: {
+                    resolved: (_data, _notification) => {
+                      setIsRescanning(false);
+                      return t('maintenance.RescanImageFinished');
+                    },
+                  },
+                },
+              };
+            } else {
+              throw new Error(t('maintenance.RescanFailed'));
+            }
+          },
+          rejected: (data, _) => {
+            const error = data as Error | undefined;
+            setIsRescanning(false);
+            return error?.message || t('maintenance.RescanFailed');
+          },
+        },
       },
     });
-    // If values can be passed through resolve, please refactor the following function using promises
-    baiClient.maintenance
-      .rescan_images()
-      .then(({ rescan_images }) => {
-        if (rescan_images.ok) {
-          upsertNotification({
-            key: notiKey,
-            backgroundTask: {
-              status: 'pending',
-              percent: 0,
-              taskId: rescan_images.task_id,
-              statusDescriptions: {
-                pending: t('maintenance.RescanImageScanning'),
-                resolved: t('maintenance.RescanImageFinished'),
-                rejected: t('maintenance.RescanFailed'),
-              },
-            },
-            duration: 0,
-          });
-        } else {
-          throw new Error(t('maintenance.RescanFailed'));
-        }
-      })
-      .catch((err: any) => {
-        upsertNotification({
-          key: notiKey,
-          // description: painKiller.relieve(err?.title),
-          backgroundTask: {
-            status: 'rejected',
-            statusDescriptions: {
-              rejected: err?.message || t('maintenance.RescanFailed'),
-            },
-          },
-          open: true,
-          duration: CLOSING_DURATION,
-        });
-      })
-      .finally(() => {
-        setIsRescanning(false);
-      });
   };
 
   const settingGroupList: SettingGroup = [

react/src/components/ModelCloneModal.tsx

@@ -109,7 +109,7 @@ const ModelCloneModal: React.FC<ModelCloneModalProps> = ({
                         status: 'pending',
                         percent: 0,
                         taskId: data.bgtask_id,
-                        statusDescriptions: {
+                        onChange: {
                           pending: t('data.folders.FolderClonePending'),
                           resolved: t('data.folders.FolderCloned'),
                           rejected: t('data.folders.FolderCloneFailed'),

react/src/hooks/useBAINotification.tsx

@@ -22,17 +22,39 @@ export interface NotificationState
     taskId?: string;
     percent?: number;
     status: 'pending' | 'rejected' | 'resolved';
-    statusDescriptions?: {
-      pending?: string;
-      resolved?: string;
-      rejected?: string;
+    onChange?: {
+      pending?:
+        | string
+        | Partial<NotificationState>
+        | ((
+            data: unknown,
+            notification: NotificationStateForOnChange,
+          ) => string | Partial<NotificationState>);
+      resolved?:
+        | string
+        | Partial<NotificationState>
+        | ((
+            data: unknown,
+            notification: NotificationStateForOnChange,
+          ) => string | Partial<NotificationState>);
+      rejected?:
+        | string
+        | Partial<NotificationState>
+        | ((
+            data: unknown,
+            notification: NotificationStateForOnChange,
+          ) => string | Partial<NotificationState>);
     };
     renderDataMessage?: (message?: string) => React.ReactNode;
-    promise?: Promise<any>;
+    promise?: Promise<unknown> | null;
   };
   extraDescription?: string;
 }
 
+export type NotificationStateForOnChange = Partial<
+  Omit<NotificationState, 'key' | 'created'>
+>;
+
 export const notificationListState = atom<NotificationState[]>([]);
 
 export const CLOSING_DURATION = 4; //second
@@ -70,30 +92,37 @@ export const useBAINotificationEffect = () => {
       ) {
         listeningPromiseKeysRef.current.push(notification.key);
         notification.backgroundTask?.promise
-          .then(() => {
-            upsertNotification({
-              key: notification.key,
-              // message: notification.message,
-              description:
-                notification.backgroundTask?.statusDescriptions?.resolved,
+          .then((data) => {
+            const updatedNotification = _.merge({}, notification, {
               backgroundTask: {
                 status: 'resolved',
               },
               duration: CLOSING_DURATION,
             });
+            const overrideData = generateOverrideByStatus(
+              updatedNotification,
+              data,
+            );
+            upsertNotification(
+              _.merge({}, updatedNotification, overrideData),
+              true,
+            );
           })
           .catch((e) => {
-            upsertNotification({
-              key: notification.key,
-              description:
-                e?.message ||
-                notification.backgroundTask?.statusDescriptions?.rejected,
+            const updatedNotification = _.merge({}, notification, {
               backgroundTask: {
                 status: 'rejected',
               },
-              // extraDescription: e?.message,
               duration: CLOSING_DURATION,
             });
+            const overrideData = generateOverrideByStatus(
+              updatedNotification,
+              e,
+            );
+            upsertNotification(
+              _.merge({}, updatedNotification, overrideData),
+              true,
+            );
           })
           .finally(() => {
             listeningPromiseKeysRef.current = _.without(
@@ -259,22 +288,33 @@ export const useSetBAINotification = () => {
   );
 
   const upsertNotification = useCallback(
-    (params: Partial<Omit<NotificationState, 'created'>>) => {
+    (
+      params: Partial<Omit<NotificationState, 'created'>>,
+      skipOverrideByStatus?: boolean,
+    ) => {
       let currentKey: React.Key | undefined;
       setNotifications((prevNotifications: NotificationState[]) => {
         let nextNotifications: NotificationState[];
         const existingIndex = params.key
           ? _.findIndex(prevNotifications, { key: params.key })
           : -1;
-        const newNotification: NotificationState = _.merge(
+        const existingNotification =
+          existingIndex > -1 ? prevNotifications[existingIndex] : undefined;
+        let newNotification: NotificationState = _.merge(
           {}, // start with empty object
-          prevNotifications[existingIndex],
+          existingNotification,
           params,
           {
             key: params.key || uuidv4(),
-            created: new Date().toISOString(),
+            created: existingNotification?.created ?? new Date().toISOString(),
           },
         );
+
+        if (!skipOverrideByStatus) {
+          const overrideData = generateOverrideByStatus(newNotification);
+          newNotification = _.merge({}, newNotification, overrideData);
+        }
+
         // This is to check if the notification should be updated using ant.d notification
         const shouldUpdateUsingAPI =
           (_.isEmpty(params.key) && params.open) ||
@@ -286,12 +326,6 @@ export const useSetBAINotification = () => {
             !_activeNotificationKeys.includes(newNotification.key) &&
             params.open);
 
-        // override description according to background task status
-        newNotification.description =
-          newNotification.backgroundTask?.statusDescriptions?.[
-            newNotification.backgroundTask?.status
-          ] || newNotification.description;
-
         if (existingIndex >= 0) {
           nextNotifications = [
             ...prevNotifications.slice(0, existingIndex),
@@ -371,3 +405,38 @@ export const useSetBAINotification = () => {
     destroyAllNotifications,
   };
 };
+
+function generateOverrideByStatus(
+  notification?: NotificationState,
+  dataOrError?: any,
+) {
+  const currentHandler =
+    notification?.backgroundTask?.onChange?.[
+      notification.backgroundTask?.status
+    ];
+
+  if (currentHandler) {
+    const overrideData = _.isFunction(currentHandler)
+      ? currentHandler(dataOrError, notification)
+      : currentHandler;
+    if (typeof overrideData === 'string') {
+      return {
+        description: overrideData,
+      };
+    } else {
+      return overrideData;
+    }
+  } else {
+    // If there is no handler for rejected case, set description using error message
+    if (
+      notification?.backgroundTask?.status === 'rejected' &&
+      dataOrError?.message
+    ) {
+      return {
+        // TODO: need to sanitize the error message using Painkiller
+        description: dataOrError.message,
+      };
+    }
+  }
+  return {};
+}

react/src/pages/SessionLauncherPage.tsx

@@ -570,7 +570,7 @@ const SessionLauncherPage = () => {
           backgroundTask: {
             promise: Promise.all(sessionPromises),
             status: 'pending',
-            statusDescriptions: {
+            onChange: {
               pending: t('session.PreparingSession'),
               resolved: t('eduapi.ComputeSessionPrepared'),
             },