fix: NotificationServicesController - check keyring state when using stateChange event

packages/notification-services-controller/CHANGELOG.md

@@ -11,6 +11,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Bump `@metamask/base-controller` from ^8.0.0 to ^8.0.1 ([#5722](https://github.com/MetaMask/core/pull/5722))
 
+### Fixed
+
+- add a check inside the `KeyringController:stateChange` subscription inside `NotificationServicesController` to prevent infinite updates ([#5731](https://github.com/MetaMask/core/pull/5731))
+  - As we invoke a `KeyringController:withKeyring` inside the `KeyringController:stateChange` event subscription,
+    we are causing many infinite updates which block other controllers from performing state updates.
+  - We now check the size of keyrings from the `KeyringController:stateChange` to better assume when keyrings have been added
+
 ## [6.0.0]
 
 ### Changed

packages/notification-services-controller/src/NotificationServicesController/NotificationServicesController.test.ts

@@ -103,50 +103,30 @@ describe('metamask-notifications - init()', () => {
   const actPublishKeyringStateChange = async (
     // eslint-disable-next-line @typescript-eslint/no-explicit-any
     messenger: any,
+    accounts: string[] = ['0x111', '0x222'],
   ) => {
     messenger.publish(
       'KeyringController:stateChange',
-      {} as KeyringControllerState,
+      {
+        keyrings: [{ accounts }],
+      } as KeyringControllerState,
       [],
     );
   };
 
-  it('keyring Change Event but feature not enabled will not add or remove triggers', async () => {
-    const { messenger, globalMessenger, mockWithKeyring } = arrangeMocks();
-
-    // initialize controller with 1 address
-    mockWithKeyring.mockResolvedValueOnce([ADDRESS_1]);
-    const controller = new NotificationServicesController({
-      messenger,
-      env: { featureAnnouncements: featureAnnouncementsEnv },
-    });
-    controller.init();
-
-    const mockUpdate = jest
-      .spyOn(controller, 'updateOnChainTriggersByAccount')
-      .mockResolvedValue({} as UserStorage);
-    const mockDelete = jest
-      .spyOn(controller, 'deleteOnChainTriggersByAccount')
-      .mockResolvedValue({} as UserStorage);
-
-    // listAccounts has a new address
-    mockWithKeyring.mockResolvedValueOnce([ADDRESS_1, ADDRESS_2]);
-    await actPublishKeyringStateChange(globalMessenger);
-
-    expect(mockUpdate).not.toHaveBeenCalled();
-    expect(mockDelete).not.toHaveBeenCalled();
-  });
-
-  it('keyring Change Event with new triggers will update triggers correctly', async () => {
-    const { messenger, globalMessenger, mockWithKeyring } = arrangeMocks();
-
+  const arrangeActAssertKeyringTest = async (
+    controllerState?: Partial<NotificationServicesControllerState>,
+  ) => {
+    const mocks = arrangeMocks();
+    const { messenger, globalMessenger, mockWithKeyring } = mocks;
     // initialize controller with 1 address
     const controller = new NotificationServicesController({
       messenger,
       env: { featureAnnouncements: featureAnnouncementsEnv },
       state: {
         isNotificationServicesEnabled: true,
-        subscriptionAccountsSeen: [ADDRESS_1],
+        subscriptionAccountsSeen: [],
+        ...controllerState,
       },
     });
     controller.init();
@@ -160,7 +140,7 @@ describe('metamask-notifications - init()', () => {
 
     const act = async (addresses: string[], assertion: () => void) => {
       mockWithKeyring.mockResolvedValueOnce(addresses);
-      await actPublishKeyringStateChange(globalMessenger);
+      await actPublishKeyringStateChange(globalMessenger, addresses);
       await waitFor(() => {
         assertion();
       });
@@ -170,6 +150,26 @@ describe('metamask-notifications - init()', () => {
       mockDelete.mockClear();
     };
 
+    return { act, mockUpdate, mockDelete };
+  };
+
+  it('event KeyringController:stateChange will not add or remove triggers when feature is disabled', async () => {
+    const { act, mockUpdate, mockDelete } = await arrangeActAssertKeyringTest({
+      isNotificationServicesEnabled: false,
+    });
+
+    // listAccounts has a new address
+    await act([ADDRESS_1, ADDRESS_2], () => {
+      expect(mockUpdate).not.toHaveBeenCalled();
+      expect(mockDelete).not.toHaveBeenCalled();
+    });
+  });
+
+  it('event KeyringController:stateChange will update notification triggers when keyring accounts change', async () => {
+    const { act, mockUpdate, mockDelete } = await arrangeActAssertKeyringTest({
+      subscriptionAccountsSeen: [ADDRESS_1],
+    });
+
     // Act - if list accounts has been seen, then will not update
     await act([ADDRESS_1], () => {
       expect(mockUpdate).not.toHaveBeenCalled();
@@ -195,6 +195,22 @@ describe('metamask-notifications - init()', () => {
     });
   });
 
+  it('event KeyringController:stateChange will update only once when if the number of keyring accounts do not change', async () => {
+    const { act, mockUpdate, mockDelete } = await arrangeActAssertKeyringTest();
+
+    // Act - First list of items, so will update
+    await act([ADDRESS_1, ADDRESS_2], () => {
+      expect(mockUpdate).toHaveBeenCalled();
+      expect(mockDelete).not.toHaveBeenCalled();
+    });
+
+    // Act - Since number of addresses in keyring has not changed, will not update
+    await act([ADDRESS_1, ADDRESS_2], () => {
+      expect(mockUpdate).not.toHaveBeenCalled();
+      expect(mockDelete).not.toHaveBeenCalled();
+    });
+  });
+
   const arrangeActInitialisePushNotifications = (
     modifications?: (mocks: ReturnType<typeof arrangeMocks>) => void,
   ) => {

packages/notification-services-controller/src/NotificationServicesController/NotificationServicesController.ts

@@ -16,6 +16,7 @@ import {
   type KeyringControllerUnlockEvent,
   type KeyringControllerWithKeyringAction,
   KeyringTypes,
+  type KeyringControllerState,
 } from '@metamask/keyring-controller';
 import type {
   AuthenticationController,
@@ -501,8 +502,12 @@ export default class NotificationServicesController extends BaseController<
     subscribe: () => {
       this.messagingSystem.subscribe(
         'KeyringController:stateChange',
-        async () => {
-          if (!this.state.isNotificationServicesEnabled) {
+        async (totalAccounts, prevTotalAccounts) => {
+          const hasTotalAccountsChanged = totalAccounts !== prevTotalAccounts;
+          if (
+            !this.state.isNotificationServicesEnabled ||
+            !hasTotalAccountsChanged
+          ) {
             return;
           }
 
@@ -518,6 +523,12 @@ export default class NotificationServicesController extends BaseController<
           }
           await Promise.all(promises);
         },
+        (state: KeyringControllerState) => {
+          return (
+            state?.keyrings?.flatMap?.((keyring) => keyring.accounts)?.length ??
+            0
+          );
+        },
       );
     },
   };