feat(wallet): Group ledger accounts by the device main public key

apps/wallet/src/background/accounts/index.ts

@@ -278,7 +278,11 @@ export async function accountsHandleUIMessage(msg: Message, uiConnection: UiConn
             const { password, accounts } = payload.args;
             for (const aLedgerAccount of accounts) {
                 newSerializedAccounts.push(
-                    await LedgerAccount.createNew({ ...aLedgerAccount, password }),
+                    await LedgerAccount.createNew({
+                        ...aLedgerAccount,
+                        password,
+                        mainPublicKey: payload.args.mainPublicKey,
+                    }),
                 );
             }
         } else if (type === AccountType.KeystoneDerived) {

apps/wallet/src/background/accounts/ledgerAccount.ts

@@ -14,6 +14,7 @@ import {
 
 export interface LedgerAccountSerialized extends SerializedAccount {
     type: AccountType.LedgerDerived;
+    mainPublicKey?: string;
     derivationPath: string;
     // just used for authentication nothing is stored here at the moment
     encrypted: string;
@@ -22,6 +23,7 @@ export interface LedgerAccountSerialized extends SerializedAccount {
 export interface LedgerAccountSerializedUI extends SerializedUIAccount {
     type: AccountType.LedgerDerived;
     derivationPath: string;
+    mainPublicKey?: string;
 }
 
 export function isLedgerAccountSerializedUI(
@@ -45,18 +47,21 @@ export class LedgerAccount
         publicKey,
         password,
         derivationPath,
+        mainPublicKey,
     }: {
         address: string;
         publicKey: string | null;
         password: string;
         derivationPath: string;
+        mainPublicKey?: string;
     }): Promise<Omit<LedgerAccountSerialized, 'id'>> {
         return {
             type: AccountType.LedgerDerived,
             address,
             publicKey,
             encrypted: await encrypt(password, {}),
             derivationPath,
+            mainPublicKey,
             lastUnlockedOn: null,
             selected: false,
             nickname: null,
@@ -97,7 +102,7 @@ export class LedgerAccount
     }
 
     async toUISerialized(): Promise<LedgerAccountSerializedUI> {
-        const { address, type, publicKey, derivationPath, selected, nickname } =
+        const { address, type, publicKey, derivationPath, selected, nickname, mainPublicKey } =
             await this.getStoredData();
         return {
             id: this.id,
@@ -106,6 +111,7 @@ export class LedgerAccount
             isLocked: await this.isLocked(),
             publicKey,
             derivationPath,
+            mainPublicKey,
             lastUnlockedOn: await this.lastUnlockedOn,
             selected,
             nickname,

apps/wallet/src/background/connections/uiConnection.ts

@@ -316,6 +316,7 @@ export class UiConnection extends Connection {
                                 await LedgerAccount.createNew({
                                     password: payload.sourceStrategy.password,
                                     ...address,
+                                    mainPublicKey: payload.sourceStrategy.mainPublicKey,
                                 }),
                             );
                         }

apps/wallet/src/shared/messaging/messages/payloads/accounts-finder/types.ts

@@ -22,6 +22,7 @@ export type SourceStrategyToPersist =
     | {
           type: 'ledger';
           password: string;
+          mainPublicKey: string;
           addresses: {
               address: string;
               derivationPath: string;

apps/wallet/src/shared/messaging/messages/payloads/methodPayload.ts

@@ -55,6 +55,7 @@ type MethodPayloads = {
         | { type: AccountType.PrivateKeyDerived; keyPair: string; password: string }
         | {
               type: AccountType.LedgerDerived;
+              mainPublicKey: string;
               accounts: { publicKey: string; derivationPath: string; address: string }[];
               password: string;
           }

apps/wallet/src/ui/app/components/accounts/AccountsFormContext.tsx

@@ -32,6 +32,7 @@ export type AccountsFormValues =
     | { type: AccountsFormType.ImportPrivateKey; keyPair: string }
     | {
           type: AccountsFormType.ImportLedger;
+          mainPublicKey: string;
           accounts: { publicKey: string; derivationPath: string; address: string }[];
       }
     | {

apps/wallet/src/ui/app/components/ledger/useDeriveLedgerAccounts.ts

@@ -14,7 +14,10 @@ type LedgerAccountKeys = 'address' | 'publicKey' | 'type' | 'derivationPath';
 
 export type DerivedLedgerAccount = Pick<LedgerAccountSerializedUI, LedgerAccountKeys>;
 interface UseDeriveLedgerAccountOptions
-    extends Pick<UseQueryOptions<DerivedLedgerAccount[], unknown>, 'select'> {
+    extends Pick<
+        UseQueryOptions<{ accounts: DerivedLedgerAccount[]; mainPublicKey: string }, unknown>,
+        'select'
+    > {
     numAccountsToDerive: number;
 }
 
@@ -40,22 +43,28 @@ async function deriveAccountsFromLedger(
     iotaLedgerClient: IotaLedgerClient,
     numAccountsToDerive: number,
 ) {
-    const ledgerAccounts: DerivedLedgerAccount[] = [];
+    const accounts: DerivedLedgerAccount[] = [];
     const derivationPaths = getDerivationPathsForLedger(numAccountsToDerive);
 
+    const mainPublicKeyResult = await iotaLedgerClient.getPublicKey(`m/44'/4218'/0'/0'/0'`);
+    const mainPublicKey = new Ed25519PublicKey(mainPublicKeyResult.publicKey);
+
     for (const derivationPath of derivationPaths) {
         const publicKeyResult = await iotaLedgerClient.getPublicKey(derivationPath);
         const publicKey = new Ed25519PublicKey(publicKeyResult.publicKey);
         const iotaAddress = publicKey.toIotaAddress();
-        ledgerAccounts.push({
+        accounts.push({
             type: AccountType.LedgerDerived,
             address: iotaAddress,
             derivationPath,
             publicKey: publicKey.toBase64(),
         });
     }
 
-    return ledgerAccounts;
+    return {
+        mainPublicKey: mainPublicKey.toBase64(),
+        accounts,
+    };
 }
 
 function getDerivationPathsForLedger(numDerivations: number) {

apps/wallet/src/ui/app/helpers/accounts.ts

@@ -4,10 +4,19 @@
 
 import { AccountType, type SerializedUIAccount } from '_src/background/accounts/account';
 import { isKeystoneAccountSerializedUI } from '_src/background/accounts/keystoneAccount';
+import { isLedgerAccountSerializedUI } from '_src/background/accounts/ledgerAccount';
 import { isMnemonicSerializedUiAccount } from '_src/background/accounts/mnemonicAccount';
 import { isSeedSerializedUiAccount } from '_src/background/accounts/seedAccount';
 
 export function getKey(account: SerializedUIAccount): string {
+    if (isMnemonicSerializedUiAccount(account)) return account.sourceID;
+    if (isSeedSerializedUiAccount(account)) return account.sourceID;
+    if (isKeystoneAccountSerializedUI(account)) return account.sourceID;
+    if (isLedgerAccountSerializedUI(account) && account.mainPublicKey) return account.mainPublicKey;
+    return account.type;
+}
+
+export function getSourceId(account: SerializedUIAccount): string {
     if (isMnemonicSerializedUiAccount(account)) return account.sourceID;
     if (isSeedSerializedUiAccount(account)) return account.sourceID;
     if (isKeystoneAccountSerializedUI(account)) return account.sourceID;
@@ -27,15 +36,19 @@ export function groupByType(accounts: SerializedUIAccount[]) {
         (acc, account) => {
             const byType = acc[account.type] || (acc[account.type] = {});
             const key = getKey(account);
-            (byType[key] || (byType[key] = [])).push(account);
+            const sourceId = getSourceId(account);
+            (byType[key] || (byType[key] = { sourceId, accounts: [] })).accounts.push(account);
             return acc;
         },
         DEFAULT_SORT_ORDER.reduce(
             (acc, type) => {
                 acc[type] = {};
                 return acc;
             },
-            {} as Record<AccountType, Record<string, SerializedUIAccount[]>>,
+            {} as Record<
+                AccountType,
+                Record<string, { sourceId: string; accounts: SerializedUIAccount[] }>
+            >,
         ),
     );
 }

apps/wallet/src/ui/app/hooks/useAccountGroups.ts

@@ -17,7 +17,9 @@ export function useAccountGroups() {
     const list = () => {
         return DEFAULT_SORT_ORDER.flatMap((type) => {
             const group = sortedAndGroupedAccounts[type];
-            return Object.values(group).flat();
+            return Object.values(group)
+                .map(({ accounts }) => accounts)
+                .flat();
         });
     };
 

apps/wallet/src/ui/app/hooks/useAccountsFinder.ts

@@ -92,6 +92,10 @@ export function useAccountsFinder({
         let sourceStrategyToPersist: SourceStrategyToPersist | undefined = undefined;
 
         if (sourceStrategy.type == 'ledger') {
+            const client = ledgerIotaClient.iotaLedgerClient!;
+            const mainPublicKeyResult = await client.getPublicKey(`m/44'/4218'/0'/0'/0'`);
+            const mainPublicKey = new Ed25519PublicKey(mainPublicKeyResult.publicKey);
+
             const addresses = await Promise.all(
                 foundAddresses.map(async (address) => {
                     const derivationPath = makeDerivationPath(address.bipPath);
@@ -107,6 +111,7 @@ export function useAccountsFinder({
             sourceStrategyToPersist = {
                 ...sourceStrategy,
                 addresses,
+                mainPublicKey: mainPublicKey.toBase64(),
             };
         } else {
             const bipPaths = foundAddresses.map((address) => address.bipPath);

apps/wallet/src/ui/app/hooks/useCreateAccountsMutation.ts

@@ -139,6 +139,7 @@ export function useCreateAccountsMutation() {
                     type: AccountType.LedgerDerived,
                     accounts: accountsFormValues.accounts,
                     password: password!,
+                    mainPublicKey: accountsFormValues.mainPublicKey,
                 });
             } else if (
                 type === AccountsFormType.ImportKeystone &&

apps/wallet/src/ui/app/pages/accounts/ImportLedgerAccountsPage.tsx

@@ -33,10 +33,14 @@ export function ImportLedgerAccountsPage() {
         isError: encounteredDerviceAccountsError,
     } = useDeriveLedgerAccounts({
         numAccountsToDerive: NUM_LEDGER_ACCOUNTS_TO_DERIVE_BY_DEFAULT,
-        select: (ledgerAccounts) => {
-            return ledgerAccounts.filter(
-                ({ address }) => !existingAccounts?.some((account) => account.address === address),
-            );
+        select: ({ accounts, mainPublicKey }) => {
+            return {
+                accounts: accounts.filter(
+                    ({ address }) =>
+                        !existingAccounts?.some((account) => account.address === address),
+                ),
+                mainPublicKey,
+            };
         },
     });
 
@@ -61,7 +65,7 @@ export function ImportLedgerAccountsPage() {
         },
         [setSelectedLedgerAccounts],
     );
-    const numImportableAccounts = ledgerAccounts?.length;
+    const numImportableAccounts = ledgerAccounts?.accounts?.length;
     const numSelectedAccounts = selectedLedgerAccounts.size;
     const areAllAccountsImported = numImportableAccounts === 0;
     const isUnlockButtonDisabled = numSelectedAccounts === 0;
@@ -76,7 +80,7 @@ export function ImportLedgerAccountsPage() {
         importLedgerAccountsBody = (
             <div className="max-h-[530px] w-full overflow-auto">
                 <AccountList
-                    accounts={ledgerAccounts}
+                    accounts={ledgerAccounts.accounts}
                     selectedAccounts={selectedLedgerAccounts}
                     onAccountClick={onAccountClick}
                     selectAll={selectAllAccounts}
@@ -88,17 +92,21 @@ export function ImportLedgerAccountsPage() {
     function selectAllAccounts() {
         const areAllAccountsSelected = numSelectedAccounts === numImportableAccounts;
         if (ledgerAccounts && !areAllAccountsSelected) {
-            setSelectedLedgerAccounts(new Set(ledgerAccounts.map((acc) => acc.address)));
+            setSelectedLedgerAccounts(new Set(ledgerAccounts.accounts.map((acc) => acc.address)));
         } else if (areAllAccountsSelected) {
             setSelectedLedgerAccounts(new Set());
         }
     }
 
     function handleNextClick() {
+        if (!ledgerAccounts) {
+            return;
+        }
         setAccountsFormValues({
             type: AccountsFormType.ImportLedger,
+            mainPublicKey: ledgerAccounts.mainPublicKey,
             accounts:
-                ledgerAccounts
+                ledgerAccounts.accounts
                     ?.filter((acc) => selectedLedgerAccounts.has(acc.address))
                     .map(({ address, derivationPath, publicKey }) => ({
                         address,

apps/wallet/src/ui/app/pages/accounts/forgot-password/RecoverManyPage.tsx

@@ -57,7 +57,7 @@ export function RecoverManyPage() {
                 </span>
                 <div className="flex w-full flex-1 flex-col gap-lg overflow-auto">
                     {mnemonicAccounts.length > 0
-                        ? mnemonicAccounts.map(([sourceID, accounts], index) => {
+                        ? mnemonicAccounts.map(([sourceID, { accounts }], index) => {
                               const recoveryData = value.find(
                                   ({ accountSourceID }) => accountSourceID === sourceID,
                               );
@@ -81,7 +81,7 @@ export function RecoverManyPage() {
                           })
                         : null}
                     {seedAccounts.length > 0
-                        ? seedAccounts.map(([sourceID, accounts], index) => {
+                        ? seedAccounts.map(([sourceID, { accounts }], index) => {
                               const recoveryData = value.find(
                                   ({ accountSourceID }) => accountSourceID === sourceID,
                               );

apps/wallet/src/ui/app/pages/accounts/manage/ManageAccountsPage.tsx

@@ -30,18 +30,22 @@ export function ManageAccountsPage() {
                 <div className="flex flex-1 flex-col overflow-y-auto">
                     <div ref={outerRef} className="relative">
                         {Object.entries(groupedAccounts).map(([type, accountGroups]) =>
-                            Object.entries(accountGroups).map(([key, accounts], index) => {
-                                return (
-                                    <AccountGroup
-                                        outerRef={outerRef}
-                                        key={`${type}-${key}`}
-                                        accounts={accounts}
-                                        accountSourceID={key}
-                                        type={type as AccountType}
-                                        isLast={index === Object.entries(accountGroups).length - 1}
-                                    />
-                                );
-                            }),
+                            Object.entries(accountGroups).map(
+                                ([key, { sourceId, accounts }], index) => {
+                                    return (
+                                        <AccountGroup
+                                            outerRef={outerRef}
+                                            key={`${type}-${key}`}
+                                            accounts={accounts}
+                                            accountSourceID={sourceId}
+                                            type={type as AccountType}
+                                            isLast={
+                                                index === Object.entries(accountGroups).length - 1
+                                            }
+                                        />
+                                    );
+                                },
+                            ),
                         )}
                         <div id="manage-account-item-portal-container"></div>
                     </div>