[notifications] Fix: #4959 transition on close

packages/toolpad-core/src/useNotifications/NotificationsContext.ts

@@ -1,6 +1,6 @@
 'use client';
 import * as React from 'react';
-import type { ShowNotification, CloseNotification } from './useNotifications';
+import type { ShowNotification, CloseNotification, RemoveNotification } from './useNotifications';
 
 /**
  * @ignore - internal component.
@@ -9,6 +9,7 @@ import type { ShowNotification, CloseNotification } from './useNotifications';
 export interface NotificationsContextValue {
   show: ShowNotification;
   close: CloseNotification;
+  remove: RemoveNotification;
 }
 
 export const NotificationsContext = React.createContext<NotificationsContextValue | null>(null);

packages/toolpad-core/src/useNotifications/NotificationsProvider.tsx

@@ -19,6 +19,7 @@ import type {
   CloseNotification,
   ShowNotification,
   ShowNotificationOptions,
+  RemoveNotification,
 } from './useNotifications';
 import { useLocaleText, type LocaleText } from '../AppProvider/LocalizationProvider';
 
@@ -55,8 +56,7 @@ interface NotificationProps {
 function Notification({ notificationKey, open, message, options, badge }: NotificationProps) {
   const globalLocaleText = useLocaleText();
   const localeText = { ...defaultLocaleText, ...globalLocaleText };
-  const { close } = useNonNullableContext(NotificationsContext);
-
+  const { close, remove } = useNonNullableContext(NotificationsContext);
   const { severity, actionText, onAction, autoHideDuration } = options;
 
   const handleClose = React.useCallback(
@@ -69,6 +69,10 @@ function Notification({ notificationKey, open, message, options, badge }: Notifi
     [notificationKey, close],
   );
 
+  const handleExited = React.useCallback(() => {
+    remove(notificationKey);
+  }, [notificationKey, remove]);
+
   const action = (
     <React.Fragment>
       {onAction ? (
@@ -90,10 +94,25 @@ function Notification({ notificationKey, open, message, options, badge }: Notifi
 
   const props = React.useContext(RootPropsContext);
   const SnackbarComponent = props?.slots?.snackbar ?? Snackbar;
+
+  // Passing `onExited` through `externalSlotProps` here.
+  // Passing it through `additionalProps` causes it to be overwritten when
+  // transition slotProps are specified in RootPropsContext.
+  const externalSnackbarSlotProps = props?.slotProps?.snackbar?.slotProps;
+  const externalTransitionProps = externalSnackbarSlotProps?.transition;
   const snackbarSlotProps = useSlotProps({
     elementType: SnackbarComponent,
     ownerState: props,
-    externalSlotProps: props?.slotProps?.snackbar,
+    externalSlotProps: {
+      ...props?.slotProps?.snackbar,
+      slotProps: {
+        ...externalSnackbarSlotProps,
+        transition: {
+          ...externalTransitionProps,
+          onExited: handleExited,
+        },
+      },
+    },
     additionalProps: {
       open,
       autoHideDuration,
@@ -193,11 +212,18 @@ function NotificationsProvider(props: NotificationsProviderProps) {
   const close = React.useCallback<CloseNotification>((key) => {
     setState((prev) => ({
       ...prev,
-      queue: prev.queue.filter((n) => n.notificationKey !== key),
+      queue: prev.queue.map((n) => (n.notificationKey === key ? { ...n, open: false } : n)),
+    }));
+  }, []);
+
+  const remove = React.useCallback<RemoveNotification>((key) => {
+    setState((prev) => ({
+      ...prev,
+      queue: prev.queue.filter((n) => key !== n.notificationKey),
     }));
   }, []);
 
-  const contextValue = React.useMemo(() => ({ show, close }), [show, close]);
+  const contextValue = React.useMemo(() => ({ show, close, remove }), [show, close, remove]);
 
   return (
     <RootPropsContext.Provider value={props}>

packages/toolpad-core/src/useNotifications/useNotifications.test.tsx

@@ -4,7 +4,7 @@
 
 import * as React from 'react';
 import { describe, test, expect } from 'vitest';
-import { renderHook, within, screen } from '@testing-library/react';
+import { renderHook, within, screen, waitFor } from '@testing-library/react';
 import { userEvent } from '@testing-library/user-event';
 import { useNotifications } from './useNotifications';
 import { NotificationsProvider } from './NotificationsProvider';
@@ -35,6 +35,8 @@ describe('useNotifications', () => {
 
     rerender();
 
-    expect(screen.queryByRole('alert')).toBeNull();
+    await waitFor(() => {
+      expect(screen.queryByRole('alert')).toBeNull();
+    });
   });
 });

packages/toolpad-core/src/useNotifications/useNotifications.tsx

@@ -46,6 +46,15 @@ export interface CloseNotification {
   (key: string): void;
 }
 
+export interface RemoveNotification {
+  /**
+   * Remove a snackbar from the application state (after it has been closed).
+   *
+   * @param key The key of the notification to remove.
+   */
+  (key: string): void;
+}
+
 interface UseNotifications {
   show: ShowNotification;
   close: CloseNotification;