feat(snap-account-service): forward selected account group accounts

README.md

@@ -509,6 +509,7 @@ linkStyle default opacity:0.5
   signature_controller --> logging_controller;
   signature_controller --> messenger;
   signature_controller --> network_controller;
+  snap_account_service --> account_tree_controller;
   snap_account_service --> keyring_controller;
   snap_account_service --> messenger;
   social_controllers --> base_controller;

packages/snap-account-service/CHANGELOG.md

@@ -28,6 +28,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   - The service messenger now requires the `KeyringController:withController` action.
 - Add `handleKeyringSnapMessage` ([#8758](https://github.com/MetaMask/core/pull/8758))
   - This will be the new entry point for consumer that needs to forward keyring events to a account management Snap (instead of using the legacy Snap keyring instance directly).
+- Forward selected account group accounts ([#8763](https://github.com/MetaMask/core/pull/8763))
+  - This logic used to live on the clients.
+  - The service messenger now requires the `KeyringController:unlock`, `AccountTreeController:selectedAccountGroupChange` events and `AccountTreeController:getAccountGroupObject` action.
 
 ### Changed
 

packages/snap-account-service/package.json

@@ -53,6 +53,8 @@
     "test:watch": "NODE_OPTIONS=--experimental-vm-modules jest --watch"
   },
   "dependencies": {
+    "@metamask/account-api": "^1.0.4",
+    "@metamask/account-tree-controller": "^7.3.0",
     "@metamask/eth-snap-keyring": "^22.0.1",
     "@metamask/keyring-controller": "^25.5.0",
     "@metamask/messenger": "^1.2.0",

packages/snap-account-service/src/SnapAccountService.test.ts

@@ -1,3 +1,4 @@
+import type { AccountGroupId } from '@metamask/account-api';
 import type { SnapKeyring, SnapMessage } from '@metamask/eth-snap-keyring';
 import {
   KeyringControllerState,
@@ -22,6 +23,7 @@ import type {
   SnapAccountServiceOptions,
 } from './SnapAccountService';
 import { SnapAccountService } from './SnapAccountService';
+import type { AccountGroupObject } from './types';
 
 type RootMessenger = Messenger<
   MockAnyNamespace,
@@ -49,6 +51,13 @@ type Mocks = {
     getState: jest.MockedFunction<() => { keyrings: { type: string }[] }>;
     withController: jest.Mock;
   };
+  // eslint-disable-next-line @typescript-eslint/naming-convention
+  AccountTreeController: {
+    getAccountGroupObject: jest.MockedFunction<
+      (groupId: AccountGroupId) => AccountGroupObject | undefined
+    >;
+    getSelectedAccountGroup: jest.MockedFunction<() => AccountGroupId | ''>;
+  };
 };
 
 /**
@@ -82,6 +91,8 @@ function getMessenger(
       'SnapController:getRunnableSnaps',
       'KeyringController:getState',
       'KeyringController:withController',
+      'AccountTreeController:getAccountGroupObject',
+      'AccountTreeController:getSelectedAccountGroup',
     ],
     events: [
       'SnapController:stateChange',
@@ -92,6 +103,8 @@ function getMessenger(
       'SnapController:snapUnblocked',
       'SnapController:snapUninstalled',
       'KeyringController:stateChange',
+      'KeyringController:unlock',
+      'AccountTreeController:selectedAccountGroupChange',
     ],
   });
   return messenger;
@@ -131,6 +144,45 @@ function publishKeyrings(
   );
 }
 
+/**
+ * Flushes pending microtasks so that chained `await`s in fire-and-forget
+ * handlers resolve before assertions run.
+ *
+ * @returns A promise that resolves once the event loop has drained.
+ */
+async function flushMicrotasks(): Promise<void> {
+  await new Promise<void>((resolve) => setImmediate(resolve));
+}
+
+/**
+ * Publishes an AccountTreeController selectedAccountGroupChange event on the
+ * root messenger.
+ *
+ * @param rootMessenger - The root messenger.
+ * @param next - The newly selected account group ID (or '').
+ * @param previous - The previously selected account group ID (or '').
+ */
+function publishSelectedAccountGroupChange(
+  rootMessenger: RootMessenger,
+  next: AccountGroupId | '',
+  previous: AccountGroupId | '' = '',
+): void {
+  rootMessenger.publish(
+    'AccountTreeController:selectedAccountGroupChange',
+    next,
+    previous,
+  );
+}
+
+/**
+ * Publishes a KeyringController unlock event on the root messenger.
+ *
+ * @param rootMessenger - The root messenger.
+ */
+function publishUnlock(rootMessenger: RootMessenger): void {
+  rootMessenger.publish('KeyringController:unlock');
+}
+
 /**
  * Builds a minimal `TruncatedSnap` for tests.
  *
@@ -147,6 +199,16 @@ function buildSnap(id: string, hasKeyring: boolean): TruncatedSnap {
   } as MockTruncatedSnap as TruncatedSnap;
 }
 
+/**
+ * Builds a minimal `AccountGroupObject` for tests.
+ *
+ * @param accounts - The list of account IDs in the group.
+ * @returns A minimal `AccountGroupObject`.
+ */
+function buildGroup(accounts: string[]): AccountGroupObject {
+  return { accounts } as unknown as AccountGroupObject;
+}
+
 /**
  * Builds a fake {@link KeyringEntry} with the given type.
  *
@@ -199,20 +261,26 @@ function mockWithController(
  * @param mocks - The mocks object from {@link setup}.
  * @param keyring - The mocked Snap keyring methods.
  * @param keyring.handleKeyringSnapMessage - The mocked implementation.
+ * @param keyring.setSelectedAccounts - The mocked implementation.
  */
 function mockLegacySnapKeyring(
   mocks: Mocks,
   {
     handleKeyringSnapMessage,
+    setSelectedAccounts,
   }: {
-    handleKeyringSnapMessage: jest.MockedFunction<
+    handleKeyringSnapMessage?: jest.MockedFunction<
       SnapKeyring['handleKeyringSnapMessage']
     >;
+    setSelectedAccounts?: jest.MockedFunction<
+      SnapKeyring['setSelectedAccounts']
+    >;
   },
 ): void {
   const snapKeyring = {
     type: KeyringTypes.snap,
     handleKeyringSnapMessage,
+    setSelectedAccounts,
   };
   mocks.KeyringController.withController.mockImplementation(async (operation) =>
     operation({
@@ -270,6 +338,10 @@ function setup({
       getState: jest.fn().mockReturnValue({ keyrings }),
       withController: jest.fn(),
     },
+    AccountTreeController: {
+      getAccountGroupObject: jest.fn().mockReturnValue(undefined),
+      getSelectedAccountGroup: jest.fn().mockReturnValue(''),
+    },
   };
 
   rootMessenger.registerActionHandler(
@@ -288,6 +360,14 @@ function setup({
     'KeyringController:withController',
     mocks.KeyringController.withController,
   );
+  rootMessenger.registerActionHandler(
+    'AccountTreeController:getAccountGroupObject',
+    mocks.AccountTreeController.getAccountGroupObject,
+  );
+  rootMessenger.registerActionHandler(
+    'AccountTreeController:getSelectedAccountGroup',
+    mocks.AccountTreeController.getSelectedAccountGroup,
+  );
 
   const service = new SnapAccountService({ messenger, config });
 
@@ -549,4 +629,164 @@ describe('SnapAccountService', () => {
       expect(result).toBe('pong');
     });
   });
+
+  describe('on AccountTreeController:selectedAccountGroupChange', () => {
+    const MOCK_GROUP_ID = 'keyring:01JABC/group-1' as AccountGroupId;
+    const MOCK_ACCOUNTS = [
+      '00000000-0000-0000-0000-000000000001',
+      '00000000-0000-0000-0000-000000000002',
+    ];
+
+    it('forwards the selected accounts to the Snap keyring', async () => {
+      const { service, rootMessenger, mocks } = setup();
+      const setSelectedAccounts = jest.fn().mockResolvedValue(undefined);
+      mockLegacySnapKeyring(mocks, { setSelectedAccounts });
+      mocks.AccountTreeController.getAccountGroupObject.mockReturnValue(
+        buildGroup(MOCK_ACCOUNTS),
+      );
+      expect(service).toBeDefined();
+
+      publishSelectedAccountGroupChange(rootMessenger, MOCK_GROUP_ID);
+      await flushMicrotasks();
+
+      expect(
+        mocks.AccountTreeController.getAccountGroupObject,
+      ).toHaveBeenCalledWith(MOCK_GROUP_ID);
+      expect(setSelectedAccounts).toHaveBeenCalledWith(MOCK_ACCOUNTS);
+    });
+
+    it('does nothing when the new group ID is empty', async () => {
+      const { service, rootMessenger, mocks } = setup();
+      const setSelectedAccounts = jest.fn().mockResolvedValue(undefined);
+      mockLegacySnapKeyring(mocks, { setSelectedAccounts });
+      expect(service).toBeDefined();
+
+      publishSelectedAccountGroupChange(rootMessenger, '');
+      await flushMicrotasks();
+
+      expect(
+        mocks.AccountTreeController.getAccountGroupObject,
+      ).not.toHaveBeenCalled();
+      expect(setSelectedAccounts).not.toHaveBeenCalled();
+    });
+
+    it('does nothing when the account group is not found', async () => {
+      const { service, rootMessenger, mocks } = setup();
+      const setSelectedAccounts = jest.fn().mockResolvedValue(undefined);
+      mockLegacySnapKeyring(mocks, { setSelectedAccounts });
+      mocks.AccountTreeController.getAccountGroupObject.mockReturnValue(
+        undefined,
+      );
+      expect(service).toBeDefined();
+
+      publishSelectedAccountGroupChange(rootMessenger, MOCK_GROUP_ID);
+      await flushMicrotasks();
+
+      expect(
+        mocks.AccountTreeController.getAccountGroupObject,
+      ).toHaveBeenCalledWith(MOCK_GROUP_ID);
+      expect(setSelectedAccounts).not.toHaveBeenCalled();
+    });
+
+    it('logs an error when forwarding to the Snap keyring fails', async () => {
+      const { service, rootMessenger, mocks } = setup();
+      const error = new Error('forward boom');
+      const setSelectedAccounts = jest.fn().mockRejectedValue(error);
+      mockLegacySnapKeyring(mocks, { setSelectedAccounts });
+      mocks.AccountTreeController.getAccountGroupObject.mockReturnValue(
+        buildGroup(MOCK_ACCOUNTS),
+      );
+      const consoleErrorSpy = jest
+        .spyOn(console, 'error')
+        .mockImplementation(() => undefined);
+      expect(service).toBeDefined();
+
+      publishSelectedAccountGroupChange(rootMessenger, MOCK_GROUP_ID);
+      await flushMicrotasks();
+
+      expect(setSelectedAccounts).toHaveBeenCalledWith(MOCK_ACCOUNTS);
+      expect(consoleErrorSpy).toHaveBeenCalledWith(
+        'Error handling selected account group change:',
+        error,
+      );
+
+      consoleErrorSpy.mockRestore();
+    });
+  });
+
+  describe('on KeyringController:unlock', () => {
+    const MOCK_GROUP_ID = 'keyring:01JABC/group-1' as AccountGroupId;
+    const MOCK_ACCOUNTS = [
+      '00000000-0000-0000-0000-000000000001',
+      '00000000-0000-0000-0000-000000000002',
+    ];
+
+    it('forwards the currently selected account group to the Snap keyring', async () => {
+      const { service, rootMessenger, mocks } = setup();
+      const setSelectedAccounts = jest.fn().mockResolvedValue(undefined);
+      mockLegacySnapKeyring(mocks, { setSelectedAccounts });
+      mocks.AccountTreeController.getSelectedAccountGroup.mockReturnValue(
+        MOCK_GROUP_ID,
+      );
+      mocks.AccountTreeController.getAccountGroupObject.mockReturnValue(
+        buildGroup(MOCK_ACCOUNTS),
+      );
+      expect(service).toBeDefined();
+
+      publishUnlock(rootMessenger);
+      await flushMicrotasks();
+
+      expect(
+        mocks.AccountTreeController.getSelectedAccountGroup,
+      ).toHaveBeenCalledTimes(1);
+      expect(
+        mocks.AccountTreeController.getAccountGroupObject,
+      ).toHaveBeenCalledWith(MOCK_GROUP_ID);
+      expect(setSelectedAccounts).toHaveBeenCalledWith(MOCK_ACCOUNTS);
+    });
+
+    it('does nothing when no account group is selected', async () => {
+      const { service, rootMessenger, mocks } = setup();
+      const setSelectedAccounts = jest.fn().mockResolvedValue(undefined);
+      mockLegacySnapKeyring(mocks, { setSelectedAccounts });
+      mocks.AccountTreeController.getSelectedAccountGroup.mockReturnValue('');
+      expect(service).toBeDefined();
+
+      publishUnlock(rootMessenger);
+      await flushMicrotasks();
+
+      expect(
+        mocks.AccountTreeController.getAccountGroupObject,
+      ).not.toHaveBeenCalled();
+      expect(setSelectedAccounts).not.toHaveBeenCalled();
+    });
+
+    it('logs an error when forwarding to the Snap keyring fails', async () => {
+      const { service, rootMessenger, mocks } = setup();
+      const error = new Error('forward boom');
+      const setSelectedAccounts = jest.fn().mockRejectedValue(error);
+      mockLegacySnapKeyring(mocks, { setSelectedAccounts });
+      mocks.AccountTreeController.getSelectedAccountGroup.mockReturnValue(
+        MOCK_GROUP_ID,
+      );
+      mocks.AccountTreeController.getAccountGroupObject.mockReturnValue(
+        buildGroup(MOCK_ACCOUNTS),
+      );
+      const consoleErrorSpy = jest
+        .spyOn(console, 'error')
+        .mockImplementation(() => undefined);
+      expect(service).toBeDefined();
+
+      publishUnlock(rootMessenger);
+      await flushMicrotasks();
+
+      expect(setSelectedAccounts).toHaveBeenCalledWith(MOCK_ACCOUNTS);
+      expect(consoleErrorSpy).toHaveBeenCalledWith(
+        'Error forwarding selected account group on unlock:',
+        error,
+      );
+
+      consoleErrorSpy.mockRestore();
+    });
+  });
 });