From 9e315e8c1ec9604ec3036e2df3f586fe5ed8531e Mon Sep 17 00:00:00 2001 From: gantunesr <17601467+gantunesr@users.noreply.github.com> Date: Thu, 14 May 2026 16:58:48 -0400 Subject: [PATCH 1/4] feat: add batch account repository APIs --- packages/snap/src/entities/account.ts | 17 ++ packages/snap/src/entities/snap.ts | 2 +- .../src/store/BdkAccountRepository.test.ts | 167 +++++++++++++++++ .../snap/src/store/BdkAccountRepository.ts | 171 +++++++++++++++--- 4 files changed, 334 insertions(+), 23 deletions(-) diff --git a/packages/snap/src/entities/account.ts b/packages/snap/src/entities/account.ts index f7f30af65..dc17ef789 100644 --- a/packages/snap/src/entities/account.ts +++ b/packages/snap/src/entities/account.ts @@ -262,6 +262,16 @@ export type BitcoinAccountRepository = { */ getByDerivationPath(derivationPath: string[]): Promise; + /** + * Get accounts by derivation path. + * + * @param derivationPaths - derivation paths. + * @returns the accounts or null if they do not exist, in input order + */ + getByDerivationPaths( + derivationPaths: string[][], + ): Promise<(BitcoinAccount | null)[]>; + /** * Create a new account, without persisting it. * @@ -283,6 +293,13 @@ export type BitcoinAccountRepository = { */ insert(account: BitcoinAccount): Promise; + /** + * Insert accounts. + * + * @param accounts - Bitcoin accounts. + */ + insertMany(accounts: BitcoinAccount[]): Promise; + /** * Update an account. * diff --git a/packages/snap/src/entities/snap.ts b/packages/snap/src/entities/snap.ts index 9d8e9f42d..8d150ce32 100644 --- a/packages/snap/src/entities/snap.ts +++ b/packages/snap/src/entities/snap.ts @@ -13,7 +13,7 @@ import type { Inscription } from './meta-protocols'; export type SnapState = { // accountId -> account state. This is the main state of the snap. - accounts: Record; + accounts: Record; // derivationPath -> accountId. Only needed for fast lookup. derivationPaths: Record; }; diff --git a/packages/snap/src/store/BdkAccountRepository.test.ts b/packages/snap/src/store/BdkAccountRepository.test.ts index a89e12411..da29f74ab 100644 --- a/packages/snap/src/store/BdkAccountRepository.test.ts +++ b/packages/snap/src/store/BdkAccountRepository.test.ts @@ -185,6 +185,94 @@ describe('BdkAccountRepository', () => { }); }); + describe('getByDerivationPaths', () => { + const derivationPath1 = ['m', "84'", "0'", "1'"]; + const derivationPath2 = ['m', "84'", "0'", "2'"]; + const accountState1 = { + ...mockAccountState, + derivationPath: derivationPath1, + }; + const accountState2 = { + ...mockAccountState, + derivationPath: derivationPath2, + }; + const mockAccount1 = mock({ + ...mockAccount, + id: 'some-id-1', + derivationPath: derivationPath1, + }); + const mockAccount2 = mock({ + ...mockAccount, + id: 'some-id-2', + derivationPath: derivationPath2, + }); + + it('returns accounts in derivation path order with one state read per namespace', async () => { + mockSnapClient.getState + .mockResolvedValueOnce({ + "m/84'/0'/1'": 'some-id-1', + "m/84'/0'/2'": 'some-id-2', + }) + .mockResolvedValueOnce({ + 'some-id-1': accountState1, + 'some-id-2': accountState2, + }); + (BdkAccountAdapter.load as jest.Mock) + .mockReturnValueOnce(mockAccount2) + .mockReturnValueOnce(mockAccount1); + + const result = await repo.getByDerivationPaths([ + derivationPath2, + derivationPath1, + ]); + + expect(mockSnapClient.getState).toHaveBeenCalledWith('derivationPaths'); + expect(mockSnapClient.getState).toHaveBeenCalledWith('accounts'); + expect(mockSnapClient.getState).toHaveBeenCalledTimes(2); + expect(result).toStrictEqual([mockAccount2, mockAccount1]); + expect(mockSnapClient.setState).not.toHaveBeenCalled(); + }); + + it('repairs missing derivation path indexes from account state', async () => { + mockSnapClient.getState + .mockResolvedValueOnce({ + "m/84'/0'/1'": 'some-id-1', + }) + .mockResolvedValueOnce({ + 'some-id-1': accountState1, + 'some-id-2': accountState2, + }); + (BdkAccountAdapter.load as jest.Mock) + .mockReturnValueOnce(mockAccount1) + .mockReturnValueOnce(mockAccount2); + + const result = await repo.getByDerivationPaths([ + derivationPath1, + derivationPath2, + ]); + + expect(result).toStrictEqual([mockAccount1, mockAccount2]); + expect(mockSnapClient.setState).toHaveBeenCalledWith('derivationPaths', { + "m/84'/0'/1'": 'some-id-1', + "m/84'/0'/2'": 'some-id-2', + }); + }); + + it('repairs a missing derivation path index for a single lookup', async () => { + mockSnapClient.getState.mockResolvedValueOnce({}).mockResolvedValueOnce({ + 'some-id-1': accountState1, + }); + (BdkAccountAdapter.load as jest.Mock).mockReturnValueOnce(mockAccount1); + + const result = await repo.getByDerivationPaths([derivationPath1]); + + expect(result).toStrictEqual([mockAccount1]); + expect(mockSnapClient.setState).toHaveBeenCalledWith('derivationPaths', { + "m/84'/0'/1'": 'some-id-1', + }); + }); + }); + describe('getWithSigner', () => { it('returns null if account not found', async () => { mockSnapClient.getState.mockResolvedValue(null); @@ -258,6 +346,85 @@ describe('BdkAccountRepository', () => { }); }); + describe('insertMany', () => { + it('returns an empty array when there are no accounts to insert', async () => { + const result = await repo.insertMany([]); + + expect(result).toStrictEqual([]); + expect(mockSnapClient.getState).not.toHaveBeenCalled(); + expect(mockSnapClient.setState).not.toHaveBeenCalled(); + }); + + it('throws an error if any account has no wallet data', async () => { + await expect( + repo.insertMany([ + { + ...mockAccount, + id: 'missing-wallet', + takeStaged: jest.fn().mockReturnValue(undefined), + }, + mockAccount, + ]), + ).rejects.toThrow( + 'Missing changeset data for account "missing-wallet" for insertion.', + ); + expect(mockSnapClient.setState).not.toHaveBeenCalled(); + }); + + it('inserts multiple accounts with one accounts write and one derivation path write', async () => { + const existingAccountState: AccountState = { + wallet: mockWalletData, + inscriptions: [], + derivationPath: mockDerivationPath, + }; + const account1 = mock(); + account1.id = 'some-id-1'; + account1.derivationPath = ['m', "84'", "0'", "1'"]; + const account2 = mock(); + account2.id = 'some-id-2'; + account2.derivationPath = ['m', "84'", "0'", "2'"]; + (account1.takeStaged as jest.Mock) = jest + .fn() + .mockReturnValue(mockChangeSet); + (account2.takeStaged as jest.Mock) = jest + .fn() + .mockReturnValue(mockChangeSet); + mockSnapClient.getState + .mockResolvedValueOnce({ + 'existing-id': existingAccountState, + }) + .mockResolvedValueOnce({ + "m/84'/0'/0'": 'existing-id', + }); + + const result = await repo.insertMany([account1, account2]); + + expect(result).toStrictEqual([account1, account2]); + expect(mockSnapClient.setState).toHaveBeenNthCalledWith(1, 'accounts', { + 'existing-id': existingAccountState, + 'some-id-1': { + wallet: mockWalletData, + inscriptions: [], + derivationPath: ['m', "84'", "0'", "1'"], + }, + 'some-id-2': { + wallet: mockWalletData, + inscriptions: [], + derivationPath: ['m', "84'", "0'", "2'"], + }, + }); + expect(mockSnapClient.setState).toHaveBeenNthCalledWith( + 2, + 'derivationPaths', + { + "m/84'/0'/0'": 'existing-id', + "m/84'/0'/1'": 'some-id-1', + "m/84'/0'/2'": 'some-id-2', + }, + ); + }); + }); + describe('update', () => { it('does nothing if no wallet data', async () => { await repo.update({ diff --git a/packages/snap/src/store/BdkAccountRepository.ts b/packages/snap/src/store/BdkAccountRepository.ts index 37b5d4442..3e630dcc8 100644 --- a/packages/snap/src/store/BdkAccountRepository.ts +++ b/packages/snap/src/store/BdkAccountRepository.ts @@ -34,6 +34,30 @@ function toBdkFingerprint(fingerprint: number): string { return fingerprint.toString(16).padStart(8, '0'); } +/** + * @param derivationPath - Split derivation path. + * @returns Storage key for a derivation path. + */ +function getDerivationPathKey(derivationPath: string[]): string { + return derivationPath.join('/'); +} + +/** + * @param account - Account to persist. + * @param walletData - Serialized wallet data. + * @returns Account state. + */ +function getAccountState( + account: BitcoinAccount, + walletData: ChangeSet, +): AccountState { + return { + wallet: walletData.to_json(), + inscriptions: [], + derivationPath: account.derivationPath, + }; +} + export class BdkAccountRepository implements BitcoinAccountRepository { readonly #snapClient: SnapClient; @@ -49,11 +73,7 @@ export class BdkAccountRepository implements BitcoinAccountRepository { return null; } - return BdkAccountAdapter.load( - id, - account.derivationPath, - ChangeSet.from_json(account.wallet), - ); + return this.#loadAccount(id, account); } async getAll(): Promise { @@ -64,22 +84,16 @@ export class BdkAccountRepository implements BitcoinAccountRepository { return []; } - return Object.entries(accounts) - .filter(([, account]) => account !== null) - .map(([id, account]) => - BdkAccountAdapter.load( - id, - account.derivationPath, - ChangeSet.from_json(account.wallet), - ), - ); + return Object.entries(accounts).flatMap(([id, account]) => + account ? [this.#loadAccount(id, account)] : [], + ); } async getByDerivationPath( derivationPath: string[], ): Promise { const id = await this.#snapClient.getState( - `derivationPaths.${derivationPath.join('/')}`, + `derivationPaths.${getDerivationPathKey(derivationPath)}`, ); if (!id) { return null; @@ -88,6 +102,65 @@ export class BdkAccountRepository implements BitcoinAccountRepository { return this.get(id as string); } + async getByDerivationPaths( + derivationPaths: string[][], + ): Promise<(BitcoinAccount | null)[]> { + if (derivationPaths.length === 0) { + return []; + } + + const [derivationPathIndex, accounts] = await Promise.all([ + this.#snapClient.getState('derivationPaths') as Promise< + SnapState['derivationPaths'] | null + >, + this.#snapClient.getState('accounts') as Promise< + SnapState['accounts'] | null + >, + ]); + + const accountsById = accounts ?? {}; + const existingDerivationPathIndex = derivationPathIndex ?? {}; + const accountsByDerivationPath = new Map(); + + for (const [id, account] of Object.entries(accountsById)) { + if (account) { + accountsByDerivationPath.set( + getDerivationPathKey(account.derivationPath), + [id, account], + ); + } + } + + const repairs: Record = {}; + const results = derivationPaths.map((derivationPath) => { + const pathKey = getDerivationPathKey(derivationPath); + const indexedId = existingDerivationPathIndex[pathKey]; + const indexedAccount = indexedId ? accountsById[indexedId] : null; + + if (indexedId && indexedAccount) { + return this.#loadAccount(indexedId, indexedAccount); + } + + const fallback = accountsByDerivationPath.get(pathKey); + if (!fallback) { + return null; + } + + const [id, account] = fallback; + repairs[pathKey] = id; + return this.#loadAccount(id, account); + }); + + if (Object.keys(repairs).length > 0) { + await this.#snapClient.setState('derivationPaths', { + ...existingDerivationPathIndex, + ...repairs, + }); + } + + return results; + } + async getWithSigner(id: string): Promise { const accountState = (await this.#snapClient.getState( `accounts.${id}`, @@ -148,7 +221,6 @@ export class BdkAccountRepository implements BitcoinAccountRepository { async insert(account: BitcoinAccount): Promise { const { id, derivationPath } = account; - const walletData = account.takeStaged(); if (!walletData) { throw new StorageError( @@ -158,19 +230,66 @@ export class BdkAccountRepository implements BitcoinAccountRepository { await Promise.all([ this.#snapClient.setState( - `derivationPaths.${derivationPath.join('/')}`, + `derivationPaths.${getDerivationPathKey(derivationPath)}`, id, ), - this.#snapClient.setState(`accounts.${id}`, { - wallet: walletData.to_json(), - inscriptions: [], - derivationPath, - }), + this.#snapClient.setState( + `accounts.${id}`, + getAccountState(account, walletData), + ), ]); return account; } + async insertMany(accounts: BitcoinAccount[]): Promise { + if (accounts.length === 0) { + return []; + } + + if (accounts.length === 1) { + return [await this.insert(accounts[0] as BitcoinAccount)]; + } + + const accountStateEntries: [string, AccountState][] = []; + const derivationPathEntries: [string, string][] = []; + + for (const account of accounts) { + const { id, derivationPath } = account; + const walletData = account.takeStaged(); + + if (!walletData) { + throw new StorageError( + `Missing changeset data for account "${id}" for insertion.`, + ); + } + + accountStateEntries.push([id, getAccountState(account, walletData)]); + derivationPathEntries.push([getDerivationPathKey(derivationPath), id]); + } + + const [existingAccounts, existingDerivationPaths] = await Promise.all([ + this.#snapClient.getState('accounts') as Promise< + SnapState['accounts'] | null + >, + this.#snapClient.getState('derivationPaths') as Promise< + SnapState['derivationPaths'] | null + >, + ]); + + await this.#snapClient.setState('accounts', { + ...(existingAccounts ?? {}), + ...Object.fromEntries(accountStateEntries), + }); + + await this.#snapClient.setState('derivationPaths', { + ...(existingDerivationPaths ?? {}), + ...Object.fromEntries(derivationPathEntries), + }); + + return accounts; + } + async update( account: BitcoinAccount, inscriptions?: Inscription[], @@ -246,4 +365,12 @@ export class BdkAccountRepository implements BitcoinAccountRepository { return `${txid}:${vout}`; }); } + + #loadAccount(id: string, account: AccountState): BitcoinAccount { + return BdkAccountAdapter.load( + id, + account.derivationPath, + ChangeSet.from_json(account.wallet), + ); + } } From 2fa5957ce4e5a79e7c7fe3c33ec14b932ef7e573 Mon Sep 17 00:00:00 2001 From: gantunesr <17601467+gantunesr@users.noreply.github.com> Date: Thu, 14 May 2026 19:38:58 -0400 Subject: [PATCH 2/4] chore: update CHANGELOGs --- packages/snap/CHANGELOG.md | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/packages/snap/CHANGELOG.md b/packages/snap/CHANGELOG.md index 8d6b8c715..d24a0d075 100644 --- a/packages/snap/CHANGELOG.md +++ b/packages/snap/CHANGELOG.md @@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Add + +- Add `insertMany` to repository API ([#609](https://github.com/MetaMask/snap-bitcoin-wallet/pull/609)) + ### Changed - Show a confirmation dialog before signing a PSBT from KeyringHandler and sending a transfer ([#591](https://github.com/MetaMask/snap-bitcoin-wallet/pull/591)) From b5cea79af4d7fbd80878754384c29d804560c053 Mon Sep 17 00:00:00 2001 From: gantunesr <17601467+gantunesr@users.noreply.github.com> Date: Thu, 14 May 2026 19:42:11 -0400 Subject: [PATCH 3/4] fix: fix CHANGELOG --- packages/snap/CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/snap/CHANGELOG.md b/packages/snap/CHANGELOG.md index d24a0d075..e7eebdfe5 100644 --- a/packages/snap/CHANGELOG.md +++ b/packages/snap/CHANGELOG.md @@ -7,7 +7,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] -### Add +### Added - Add `insertMany` to repository API ([#609](https://github.com/MetaMask/snap-bitcoin-wallet/pull/609)) From 4670591af2f697d5329433660a6f4fcb3c42a425 Mon Sep 17 00:00:00 2001 From: gantunesr <17601467+gantunesr@users.noreply.github.com> Date: Thu, 14 May 2026 19:50:44 -0400 Subject: [PATCH 4/4] fix: preflight batch account inserts --- packages/snap/src/entities/account.ts | 7 ++++ packages/snap/src/infra/BdkAccountAdapter.ts | 24 +++++++++++- .../src/store/BdkAccountRepository.test.ts | 39 ++++++++++++++----- .../snap/src/store/BdkAccountRepository.ts | 9 ++++- 4 files changed, 68 insertions(+), 11 deletions(-) diff --git a/packages/snap/src/entities/account.ts b/packages/snap/src/entities/account.ts index dc17ef789..c66f92951 100644 --- a/packages/snap/src/entities/account.ts +++ b/packages/snap/src/entities/account.ts @@ -123,6 +123,13 @@ export type BitcoinAccount = { */ takeStaged(): ChangeSet | undefined; + /** + * Check whether a change set exists without making it unavailable for extraction. + * + * @returns true if there is a change set + */ + hasStaged(): boolean; + /** * Returns a Transaction Builder. * diff --git a/packages/snap/src/infra/BdkAccountAdapter.ts b/packages/snap/src/infra/BdkAccountAdapter.ts index a8c8a86da..425bd4406 100644 --- a/packages/snap/src/infra/BdkAccountAdapter.ts +++ b/packages/snap/src/infra/BdkAccountAdapter.ts @@ -41,6 +41,8 @@ export class BdkAccountAdapter implements BitcoinAccount { readonly #wallet: Wallet; + #staged?: ChangeSet; + readonly #capabilities: AccountCapability[]; constructor(id: string, derivationPath: string[], wallet: Wallet) { @@ -157,7 +159,13 @@ export class BdkAccountAdapter implements BitcoinAccount { } takeStaged(): ChangeSet | undefined { - return this.#wallet.take_staged(); + const staged = this.#collectStaged(); + this.#staged = undefined; + return staged; + } + + hasStaged(): boolean { + return Boolean(this.#collectStaged()); } buildTx(): TransactionBuilder { @@ -243,4 +251,18 @@ export class BdkAccountAdapter implements BitcoinAccount { new UnconfirmedTx(tx, BigInt(lastSeen)), ]); } + + #collectStaged(): ChangeSet | undefined { + const staged = this.#wallet.take_staged(); + if (!staged) { + return this.#staged; + } + + if (this.#staged) { + staged.merge(this.#staged); + } + + this.#staged = staged; + return this.#staged; + } } diff --git a/packages/snap/src/store/BdkAccountRepository.test.ts b/packages/snap/src/store/BdkAccountRepository.test.ts index da29f74ab..4f4f6fd9d 100644 --- a/packages/snap/src/store/BdkAccountRepository.test.ts +++ b/packages/snap/src/store/BdkAccountRepository.test.ts @@ -77,6 +77,7 @@ describe('BdkAccountRepository', () => { (mockAccount.takeStaged as jest.Mock) = jest .fn() .mockReturnValue(mockChangeSet); + (mockAccount.hasStaged as jest.Mock) = jest.fn().mockReturnValue(true); (mockChangeSet.to_json as jest.Mock) = jest .fn() .mockReturnValue(mockWalletData); @@ -355,19 +356,37 @@ describe('BdkAccountRepository', () => { expect(mockSnapClient.setState).not.toHaveBeenCalled(); }); - it('throws an error if any account has no wallet data', async () => { + it('throws an error without consuming staged data if any account has no wallet data', async () => { + const accountWithWalletData = mock({ + id: 'some-id-1', + derivationPath: ['m', "84'", "0'", "1'"], + }); + (accountWithWalletData.hasStaged as jest.Mock) = jest + .fn() + .mockReturnValue(true); + (accountWithWalletData.takeStaged as jest.Mock) = jest + .fn() + .mockReturnValue(mockChangeSet); + const missingWalletAccount = mock({ + id: 'missing-wallet', + derivationPath: ['m', "84'", "0'", "2'"], + }); + (missingWalletAccount.hasStaged as jest.Mock) = jest + .fn() + .mockReturnValue(false); + (missingWalletAccount.takeStaged as jest.Mock) = jest.fn(); + await expect( - repo.insertMany([ - { - ...mockAccount, - id: 'missing-wallet', - takeStaged: jest.fn().mockReturnValue(undefined), - }, - mockAccount, - ]), + repo.insertMany([accountWithWalletData, missingWalletAccount]), ).rejects.toThrow( 'Missing changeset data for account "missing-wallet" for insertion.', ); + + expect(accountWithWalletData.hasStaged).toHaveBeenCalled(); + expect(missingWalletAccount.hasStaged).toHaveBeenCalled(); + expect(accountWithWalletData.takeStaged).not.toHaveBeenCalled(); + expect(missingWalletAccount.takeStaged).not.toHaveBeenCalled(); + expect(mockSnapClient.getState).not.toHaveBeenCalled(); expect(mockSnapClient.setState).not.toHaveBeenCalled(); }); @@ -386,9 +405,11 @@ describe('BdkAccountRepository', () => { (account1.takeStaged as jest.Mock) = jest .fn() .mockReturnValue(mockChangeSet); + (account1.hasStaged as jest.Mock) = jest.fn().mockReturnValue(true); (account2.takeStaged as jest.Mock) = jest .fn() .mockReturnValue(mockChangeSet); + (account2.hasStaged as jest.Mock) = jest.fn().mockReturnValue(true); mockSnapClient.getState .mockResolvedValueOnce({ 'existing-id': existingAccountState, diff --git a/packages/snap/src/store/BdkAccountRepository.ts b/packages/snap/src/store/BdkAccountRepository.ts index 3e630dcc8..c1ec5c260 100644 --- a/packages/snap/src/store/BdkAccountRepository.ts +++ b/packages/snap/src/store/BdkAccountRepository.ts @@ -254,6 +254,14 @@ export class BdkAccountRepository implements BitcoinAccountRepository { const accountStateEntries: [string, AccountState][] = []; const derivationPathEntries: [string, string][] = []; + for (const account of accounts) { + if (!account.hasStaged()) { + throw new StorageError( + `Missing changeset data for account "${account.id}" for insertion.`, + ); + } + } + for (const account of accounts) { const { id, derivationPath } = account; const walletData = account.takeStaged(); @@ -263,7 +271,6 @@ export class BdkAccountRepository implements BitcoinAccountRepository { `Missing changeset data for account "${id}" for insertion.`, ); } - accountStateEntries.push([id, getAccountState(account, walletData)]); derivationPathEntries.push([getDerivationPathKey(derivationPath), id]); }