@aztec/wallets 0.0.1-commit.2448fdb → 0.0.1-commit.2606882

package/src/embedded/embedded_wallet.ts

@@ -1,16 +1,32 @@
 import { type Account, NO_FROM } from '@aztec/aztec.js/account';
 import { CallAuthorizationRequest } from '@aztec/aztec.js/authorization';
-import { type InteractionWaitOptions, type SendReturn, type WaitOpts, getGasLimits } from '@aztec/aztec.js/contracts';
-import type { Aliased, SendOptions } from '@aztec/aztec.js/wallet';
+import {
+  type InteractionWaitOptions,
+  NO_WAIT,
+  type SendReturn,
+  type WaitOpts,
+  getGasLimits,
+} from '@aztec/aztec.js/contracts';
+import type {
+  Aliased,
+  ExecuteUtilityOptions,
+  PrivateEvent,
+  PrivateEventFilter,
+  ProfileOptions,
+  SendOptions,
+  SimulateOptions,
+} from '@aztec/aztec.js/wallet';
 import { AccountManager, TxSimulationResultWithAppOffset } from '@aztec/aztec.js/wallet';
 import type { DefaultAccountEntrypointOptions } from '@aztec/entrypoints/account';
 import { DefaultEntrypoint } from '@aztec/entrypoints/default';
 import { Fq, Fr } from '@aztec/foundation/curves/bn254';
 import type { Logger } from '@aztec/foundation/log';
+import type { AztecAsyncKVStore } from '@aztec/kv-store';
 import type { PXEConfig, PXECreationOptions } from '@aztec/pxe/client/lazy';
 import type { PXE } from '@aztec/pxe/server';
+import type { ContractArtifact, EventMetadataDefinition, FunctionCall } from '@aztec/stdlib/abi';
 import { AztecAddress } from '@aztec/stdlib/aztec-address';
-import { getContractInstanceFromInstantiationParams } from '@aztec/stdlib/contract';
+import { type ContractInstanceWithAddress, getContractClassFromArtifact } from '@aztec/stdlib/contract';
 import { GasSettings } from '@aztec/stdlib/gas';
 import type { AztecNode } from '@aztec/stdlib/interfaces/client';
 import { deriveSigningKey } from '@aztec/stdlib/keys';
@@ -19,7 +35,9 @@ import {
   ExecutionPayload,
   SimulationOverrides,
   type TxExecutionRequest,
+  type TxProfileResult,
   TxStatus,
+  type UtilityExecutionResult,
   collectOffchainEffects,
   mergeExecutionPayloads,
 } from '@aztec/stdlib/tx';
@@ -39,10 +57,20 @@ export function splitPxeOptions(pxe?: EmbeddedWalletPXEOptions): {
   if (!pxe) {
     return { config: {}, creation: {} };
   }
-  const { loggers, loggerActorLabel, proverOrOptions, store, simulator, ...config } = pxe;
-  return { config, creation: { loggers, loggerActorLabel, proverOrOptions, store, simulator } };
+  const { loggers, loggerActorLabel, proverOrOptions, store, simulator, hooks, preloadedContractsProvider, ...config } =
+    pxe;
+  return {
+    config,
+    creation: { loggers, loggerActorLabel, proverOrOptions, store, simulator, hooks, preloadedContractsProvider },
+  };
 }
 
+/** Options for the EmbeddedWallet's own DB (accounts, senders — distinct from PXE state). */
+export type EmbeddedWalletDBOptions = {
+  /** Override the wallet DB backend. If omitted, an IndexedDB (browser) / LMDB (node) store is created. */
+  store?: AztecAsyncKVStore;
+};
+
 export type EmbeddedWalletOptions = {
   /** Parent logger. Child loggers are derived via createChild() for each subsystem. */
   logger?: Logger;
@@ -50,6 +78,8 @@ export type EmbeddedWalletOptions = {
   ephemeral?: boolean;
   /** PXE configuration and dependency overrides (custom store, prover, simulator). */
   pxe?: EmbeddedWalletPXEOptions;
+  /** Wallet DB dependency overrides (custom store). */
+  walletDb?: EmbeddedWalletDBOptions;
   /**
    * Override PXE configuration.
    * @deprecated Use `pxe` instead.
@@ -67,6 +97,10 @@ const DEFAULT_ESTIMATED_GAS_PADDING = 0.1;
 export class EmbeddedWallet extends BaseWallet {
   protected estimatedGasPadding = DEFAULT_ESTIMATED_GAS_PADDING;
 
+  // Stub class ids, populated on wallet startup
+  // to avoid redundant work per simulation
+  protected stubClassIds = new Map<AccountType, Fr>();
+
   constructor(
     pxe: PXE,
     aztecNode: AztecNode,
@@ -118,6 +152,9 @@ export class EmbeddedWallet extends BaseWallet {
     executionPayload: ExecutionPayload,
     opts: SendOptions<W>,
   ): Promise<SendReturn<W>> {
+    // PXE has autoSync disabled by the embedded wallet entrypoints, so we sync once here to cover
+    // both the inner simulateTx (via simulateViaEntrypoint) and the proveTx that super.sendTx
+    await this.pxe.sync();
     const feeOptions = await this.completeFeeOptions({
       from: opts.from,
       feePayer: executionPayload.feePayer,
@@ -132,6 +169,7 @@ export class EmbeddedWallet extends BaseWallet {
       feeOptions,
       additionalScopes: opts.additionalScopes,
       skipTxValidation: true,
+      sendMessagesAs: opts.sendMessagesAs,
     });
 
     const offchainEffects = collectOffchainEffects(simulationResult.privateExecutionResult);
@@ -164,24 +202,92 @@ export class EmbeddedWallet extends BaseWallet {
       gasLimits: opts.fee?.gasSettings?.gasLimits ?? estimated.gasLimits,
       teardownGasLimits: opts.fee?.gasSettings?.teardownGasLimits ?? estimated.teardownGasLimits,
     });
-    const waitOpts: WaitOpts = typeof opts.wait === 'object' ? opts.wait : {};
-
-    if (!waitOpts?.waitForStatus) {
-      // Default to PROPOSED so the wait returns as soon as the tx lands in a proposed L2 block,
-      // rather than waiting until the end of the slot for the checkpoint to be published to L1.
-      // This is what makes MBPS (Multiple Blocks Per Slot) actually improve UX: with CHECKPOINTED
-      // we'd block until L1 inclusion regardless of how early in the slot the tx was sequenced.
-      // The tradeoff is a weaker guarantee — a proposed block only becomes canonical once it (or
-      // a later block in the same slot) is checkpointed, so a tx could be re-orged out if the
-      // proposer fails to publish to L1 (which should be rare, since they'd get slashed for it).
-      waitOpts!.waitForStatus = TxStatus.PROPOSED;
+    let wait: InteractionWaitOptions = opts.wait;
+    if (wait !== NO_WAIT) {
+      const callerWaitOpts: WaitOpts = typeof wait === 'object' ? wait : {};
+      wait = {
+        ...callerWaitOpts,
+        // Default to PROPOSED so the wait returns as soon as the tx lands in a proposed L2 block,
+        // rather than waiting until the end of the slot for the checkpoint to be published to L1.
+        // This is what makes MBPS (Multiple Blocks Per Slot) actually improve UX: with CHECKPOINTED
+        // we'd block until L1 inclusion regardless of how early in the slot the tx was sequenced.
+        // The tradeoff is a weaker guarantee — a proposed block only becomes canonical once it (or
+        // a later block in the same slot) is checkpointed, so a tx could be re-orged out if the
+        // proposer fails to publish to L1 (which should be rare, since they'd get slashed for it).
+        waitForStatus: callerWaitOpts.waitForStatus ?? TxStatus.PROPOSED,
+      };
     }
     return super.sendTx(executionPayload, {
       ...opts,
+      wait: wait as W,
       fee: { ...opts.fee, gasSettings },
     });
   }
 
+  /**
+   * Overrides the base simulateTx to drive PXE syncing explicitly. The PXE created by the embedded
+   * wallet has autoSync disabled (so we can share one sync across simulate+send in sendTx); for
+   * standalone simulations we still need a fresh anchor block, which we provide here.
+   */
+  public override async simulateTx(
+    executionPayload: ExecutionPayload,
+    opts: SimulateOptions,
+  ): Promise<TxSimulationResultWithAppOffset> {
+    await this.pxe.sync();
+    return super.simulateTx(executionPayload, opts);
+  }
+
+  public override async profileTx(executionPayload: ExecutionPayload, opts: ProfileOptions): Promise<TxProfileResult> {
+    await this.pxe.sync();
+    return super.profileTx(executionPayload, opts);
+  }
+
+  public override async executeUtility(
+    call: FunctionCall,
+    opts: ExecuteUtilityOptions,
+  ): Promise<UtilityExecutionResult> {
+    await this.pxe.sync();
+    return super.executeUtility(call, opts);
+  }
+
+  public override async getPrivateEvents<T>(
+    eventDef: EventMetadataDefinition,
+    eventFilter: PrivateEventFilter,
+  ): Promise<PrivateEvent<T>[]> {
+    await this.pxe.sync();
+    return super.getPrivateEvents<T>(eventDef, eventFilter);
+  }
+
+  public override async registerContract(
+    instance: ContractInstanceWithAddress,
+    artifact?: ContractArtifact,
+    secretKey?: Fr,
+  ): Promise<ContractInstanceWithAddress> {
+    // registerContract may call pxe.updateContract under the hood, which depends on a fresh anchor
+    // block to verify the current class id from the node.
+    await this.pxe.sync();
+    return super.registerContract(instance, artifact, secretKey);
+  }
+
+  /**
+   * Hashes and registers the stub class for every supported account type with PXE, populating
+   * stubClassIds. Called on wallet initialization.
+   */
+  async initStubClasses(): Promise<void> {
+    const schnorrArtifact = await this.accountContracts.getStubAccountContractArtifact('schnorr');
+    const { id: schnorrClassId } = await getContractClassFromArtifact(schnorrArtifact);
+    await this.pxe.registerContractClass(schnorrArtifact);
+
+    // ecdsa stubs share the same class id
+    const ecdsaArtifact = await this.accountContracts.getStubAccountContractArtifact('ecdsasecp256r1');
+    const { id: ecdsaClassId } = await getContractClassFromArtifact(ecdsaArtifact);
+    await this.pxe.registerContractClass(ecdsaArtifact);
+
+    this.stubClassIds.set('schnorr', schnorrClassId);
+    this.stubClassIds.set('ecdsasecp256k1', ecdsaClassId);
+    this.stubClassIds.set('ecdsasecp256r1', ecdsaClassId);
+  }
+
   /**
    * Builds contract overrides for all provided addresses by replacing their account contracts with stub implementations.
    * Uses a type-specific stub artifact so that the stub's constructor selector matches the real account's constructor.
@@ -195,7 +301,12 @@ export class EmbeddedWallet extends BaseWallet {
     for (const account of filtered) {
       const address = account.item;
       const { type } = await this.walletDB.retrieveAccount(address);
-      const stubArtifact = await this.accountContracts.getStubAccountContractArtifact(type);
+      const stubClassId = this.stubClassIds.get(type);
+      if (!stubClassId) {
+        throw new Error(
+          `Stub class for account type '${type}' was not registered at wallet init. This is a bug — initStubClasses should cover every supported AccountType.`,
+        );
+      }
 
       const originalAccount = await this.getAccountFromAddress(address);
       const completeAddress = originalAccount.getCompleteAddress();
@@ -206,15 +317,8 @@ export class EmbeddedWallet extends BaseWallet {
         );
       }
 
-      const stubConstructorArgs = type === 'schnorr' ? [Fr.ZERO, Fr.ZERO] : [Buffer.alloc(32), Buffer.alloc(32)];
-      const stubInstance = await getContractInstanceFromInstantiationParams(stubArtifact, {
-        salt: Fr.random(),
-        constructorArgs: stubConstructorArgs,
-      });
-
       contracts[address.toString()] = {
-        instance: stubInstance,
-        artifact: stubArtifact,
+        instance: { ...contractInstance, currentContractClassId: stubClassId },
       };
     }
 
@@ -230,7 +334,7 @@ export class EmbeddedWallet extends BaseWallet {
     executionPayload: ExecutionPayload,
     opts: SimulateViaEntrypointOptions,
   ): Promise<TxSimulationResultWithAppOffset> {
-    const { from, feeOptions, additionalScopes, skipTxValidation, skipFeeEnforcement } = opts;
+    const { from, feeOptions, additionalScopes, skipTxValidation, skipFeeEnforcement, sendMessagesAs } = opts;
     const scopes = this.scopesFrom(from, additionalScopes);
 
     const feeExecutionPayload = await feeOptions.walletFeePaymentMethod?.getExecutionPayload();
@@ -240,7 +344,7 @@ export class EmbeddedWallet extends BaseWallet {
     const chainInfo = await this.getChainInfo();
 
     const accountOverrides = await this.buildAccountOverrides(scopes);
-    const overrides = new SimulationOverrides(accountOverrides);
+    const overrides = new SimulationOverrides({ contracts: accountOverrides });
 
     let txRequest: TxExecutionRequest;
     if (from === NO_FROM) {
@@ -271,6 +375,7 @@ export class EmbeddedWallet extends BaseWallet {
       skipTxValidation,
       overrides,
       scopes,
+      senderForTags: this.senderForTagsFrom(from, sendMessagesAs),
     });
     const appCallOffset = await this.computeAppCallOffset(from, feeOptions);
     return TxSimulationResultWithAppOffset.fromResultAndOffset(result, appCallOffset);
@@ -301,7 +406,7 @@ export class EmbeddedWallet extends BaseWallet {
       }
     }
 
-    const accountManager = await AccountManager.create(this, secret, contract, salt);
+    const accountManager = await AccountManager.create(this, secret, contract, { salt });
 
     const instance = accountManager.getInstance();
     const existingInstance = await this.pxe.getContractInstance(instance.address);
@@ -349,7 +454,8 @@ export class EmbeddedWallet extends BaseWallet {
     this.estimatedGasPadding = value ?? DEFAULT_ESTIMATED_GAS_PADDING;
   }
 
-  stop() {
-    return this.pxe.stop();
+  async stop(): Promise<void> {
+    await this.pxe.stop();
+    await this.walletDB.close();
   }
 }

package/src/embedded/entrypoints/browser.ts

@@ -3,6 +3,8 @@ import { type Logger, createLogger } from '@aztec/foundation/log';
 import { createStore, openTmpStore } from '@aztec/kv-store/indexeddb';
 import { type PXE, type PXECreationOptions, createPXE } from '@aztec/pxe/client/lazy';
 import { type PXEConfig, getPXEConfig } from '@aztec/pxe/config';
+import { getStandardAuthRegistry } from '@aztec/standard-contracts/auth-registry/lazy';
+import { getStandardMultiCallEntrypoint } from '@aztec/standard-contracts/multi-call-entrypoint/lazy';
 
 import { LazyAccountContractsProvider } from '../account-contract-providers/lazy.js';
 import type { AccountContractsProvider } from '../account-contract-providers/types.js';
@@ -34,6 +36,7 @@ export class BrowserEmbeddedWallet extends EmbeddedWallet {
     const pxeConfig: PXEConfig = Object.assign(getPXEConfig(), {
       proverEnabled: mergedConfigOverrides.proverEnabled ?? false,
       dataDirectory: `pxe_data_${l1Contracts.rollupAddress}`,
+      autoSync: false,
       ...mergedConfigOverrides,
     });
 
@@ -43,6 +46,9 @@ export class BrowserEmbeddedWallet extends EmbeddedWallet {
 
     const pxeOptions: PXECreationOptions = {
       ...mergedCreationOverrides,
+      preloadedContractsProvider: mergedCreationOverrides.preloadedContractsProvider ?? {
+        getPreloadedContracts: async () => [await getStandardMultiCallEntrypoint(), await getStandardAuthRegistry()],
+      },
       loggers: {
         store: rootLogger.createChild('pxe:data'),
         pxe: rootLogger.createChild('pxe:service'),
@@ -53,21 +59,25 @@ export class BrowserEmbeddedWallet extends EmbeddedWallet {
 
     const pxe = await createPXE(aztecNode, pxeConfig, pxeOptions);
 
-    const walletDBStore = options.ephemeral
-      ? await openTmpStore(true)
-      : await createStore(
-          'wallet_data',
-          {
-            dataDirectory: `wallet_data_${l1Contracts.rollupAddress}`,
-            dataStoreMapSizeKb: pxeConfig.dataStoreMapSizeKb,
-            l1Contracts,
-          },
-          1,
-          rootLogger.createChild('wallet:data'),
-        );
-    const walletDB = WalletDB.init(walletDBStore, rootLogger.createChild('wallet:db').info);
+    const walletDBStore =
+      options.walletDb?.store ??
+      (options.ephemeral
+        ? await openTmpStore(true)
+        : await createStore(
+            'wallet_data',
+            {
+              dataDirectory: `wallet_data_${l1Contracts.rollupAddress}`,
+              dataStoreMapSizeKb: pxeConfig.dataStoreMapSizeKb,
+              rollupAddress: l1Contracts.rollupAddress,
+            },
+            1,
+            rootLogger.createChild('wallet:data'),
+          ));
+    const walletDB = new WalletDB(walletDBStore, rootLogger.createChild('wallet:db').info);
 
-    return new this(pxe, aztecNode, walletDB, new LazyAccountContractsProvider(), rootLogger) as T;
+    const wallet = new this(pxe, aztecNode, walletDB, new LazyAccountContractsProvider(), rootLogger) as T;
+    await wallet.initStubClasses();
+    return wallet;
   }
 }
 
@@ -75,3 +85,10 @@ export { BrowserEmbeddedWallet as EmbeddedWallet };
 export type { EmbeddedWalletOptions, EmbeddedWalletPXEOptions } from '../embedded_wallet.js';
 export { WalletDB } from '../wallet_db.js';
 export type { AccountType } from '../wallet_db.js';
+
+// At-rest encryption helpers are intentionally NOT re-exported here. They live
+// on the `@aztec/wallets/embedded/store-encryption` sub-path so consumers
+// (and bundlers) of this entrypoint don't transitively pull in
+// `@aztec/kv-store/sqlite-opfs` and its `new Worker(new URL('./worker.js'))`
+// chain into `@aztec/sqlite3mc-wasm`. Apps that don't use encryption-at-rest
+// (e.g. the playground) should never see sqlite-opfs in their bundle.

package/src/embedded/entrypoints/node.ts

@@ -3,6 +3,8 @@ import { type Logger, createLogger } from '@aztec/foundation/log';
 import { createStore, openTmpStore } from '@aztec/kv-store/lmdb-v2';
 import { type PXEConfig, getPXEConfig } from '@aztec/pxe/config';
 import { type PXE, type PXECreationOptions, createPXE } from '@aztec/pxe/server';
+import { getStandardAuthRegistry } from '@aztec/standard-contracts/auth-registry';
+import { getStandardMultiCallEntrypoint } from '@aztec/standard-contracts/multi-call-entrypoint';
 import type { AztecNode } from '@aztec/stdlib/interfaces/client';
 
 import { BundleAccountContractsProvider } from '../account-contract-providers/bundle.js';
@@ -35,6 +37,7 @@ export class NodeEmbeddedWallet extends EmbeddedWallet {
     const pxeConfig: PXEConfig = Object.assign(getPXEConfig(), {
       proverEnabled: mergedConfigOverrides.proverEnabled ?? false,
       dataDirectory: `pxe_data_${l1Contracts.rollupAddress}`,
+      autoSync: false,
       ...mergedConfigOverrides,
     });
 
@@ -44,6 +47,9 @@ export class NodeEmbeddedWallet extends EmbeddedWallet {
 
     const pxeOptions: PXECreationOptions = {
       ...mergedCreationOverrides,
+      preloadedContractsProvider: mergedCreationOverrides.preloadedContractsProvider ?? {
+        getPreloadedContracts: async () => [await getStandardMultiCallEntrypoint(), await getStandardAuthRegistry()],
+      },
       loggers: {
         store: rootLogger.createChild('pxe:data'),
         pxe: rootLogger.createChild('pxe:service'),
@@ -54,27 +60,31 @@ export class NodeEmbeddedWallet extends EmbeddedWallet {
 
     const pxe = await createPXE(aztecNode, pxeConfig, pxeOptions);
 
-    const walletDBStore = options.ephemeral
-      ? await openTmpStore(
-          `wallet_data_${l1Contracts.rollupAddress}`,
-          true,
-          undefined,
-          undefined,
-          rootLogger.createChild('wallet:data').getBindings(),
-        )
-      : await createStore(
-          'wallet_data',
-          1,
-          {
-            dataDirectory: `wallet_data_${l1Contracts.rollupAddress}`,
-            dataStoreMapSizeKb: pxeConfig.dataStoreMapSizeKb,
-            l1Contracts,
-          },
-          rootLogger.createChild('wallet:data').getBindings(),
-        );
-    const walletDB = WalletDB.init(walletDBStore, rootLogger.createChild('wallet:db').info);
+    const walletDBStore =
+      options.walletDb?.store ??
+      (options.ephemeral
+        ? await openTmpStore(
+            `wallet_data_${l1Contracts.rollupAddress}`,
+            true,
+            undefined,
+            undefined,
+            rootLogger.createChild('wallet:data').getBindings(),
+          )
+        : await createStore(
+            'wallet_data',
+            1,
+            {
+              dataDirectory: `wallet_data_${l1Contracts.rollupAddress}`,
+              dataStoreMapSizeKb: pxeConfig.dataStoreMapSizeKb,
+              rollupAddress: l1Contracts.rollupAddress,
+            },
+            rootLogger.createChild('wallet:data').getBindings(),
+          ));
+    const walletDB = new WalletDB(walletDBStore, rootLogger.createChild('wallet:db').info);
 
-    return new this(pxe, aztecNode, walletDB, new BundleAccountContractsProvider(), rootLogger) as T;
+    const wallet = new this(pxe, aztecNode, walletDB, new BundleAccountContractsProvider(), rootLogger) as T;
+    await wallet.initStubClasses();
+    return wallet;
   }
 }
 

package/src/embedded/store_encryption.ts

@@ -0,0 +1,107 @@
+/**
+ * Wallet-layer helpers for opening the embedded wallet's two encrypted stores (PXE + walletDB) as a cohesive unit.
+ *
+ * Sits on top of `@aztec/kv-store/sqlite-opfs`'s typed `SqliteEncryptionError` and adds:
+ *
+ *   - `storeName: 'pxe' | 'wallet'`, telling callers WHICH store failed.
+ *   - Cleanup: when the wallet store fails to open, ensures the already-opened PXE store is closed before the error
+ *     surfaces, so callers don't leak the SAH Pool's OPFS lock.
+ */
+import type { Logger } from '@aztec/foundation/log';
+import { AztecSQLiteOPFSStore, SqliteEncryptionError } from '@aztec/kv-store/sqlite-opfs';
+
+/** Which of the embedded wallet's two stores failed to open. */
+export type EmbeddedStoreName = 'pxe' | 'wallet';
+
+/**
+ * Thrown by {@link openEncryptedEmbeddedStores} when one of the two stores cannot be decrypted with the supplied
+ * key. The original {@link SqliteEncryptionError} is preserved as `cause`.
+ */
+export class EmbeddedWalletEncryptionError extends Error {
+  readonly storeName: EmbeddedStoreName;
+
+  constructor(storeName: EmbeddedStoreName, opts: { cause: SqliteEncryptionError }) {
+    super(`Embedded wallet '${storeName}' store could not be decrypted with the provided key`, { cause: opts.cause });
+    this.name = 'EmbeddedWalletEncryptionError';
+    this.storeName = storeName;
+  }
+}
+
+/** Configuration for {@link openEncryptedEmbeddedStores}. */
+export interface OpenEncryptedEmbeddedStoresOptions {
+  pxe: { name: string; poolDirectory?: string };
+  wallet: { name: string; poolDirectory?: string };
+}
+
+/**
+ * Internal seam for tests to inject a fake store opener. Defaults to `AztecSQLiteOPFSStore.open`. Not part of the
+ * public API.
+ *
+ * @internal
+ */
+export type OpenSqliteEncryptedStoreFn = (
+  log: Logger,
+  name: string,
+  poolDirectory: string | undefined,
+  encryptionKey: Uint8Array,
+) => Promise<AztecSQLiteOPFSStore>;
+
+const defaultOpenStore: OpenSqliteEncryptedStoreFn = (log, name, poolDirectory, encryptionKey) =>
+  AztecSQLiteOPFSStore.open(log, name, false, poolDirectory, encryptionKey);
+
+/**
+ * Opens the PXE and wallet stores in sequence, both encrypted with keys obtained from `getEncryptionKey`.
+ *
+ * The callback is invoked once per store (twice total per call) because `AztecSQLiteOPFSStore.open` *transfers*
+ * the key buffer to its worker. A single buffer would detach between the two opens.
+ *
+ * Failure modes:
+ *
+ *   - PXE store fails to decrypt → throws `EmbeddedWalletEncryptionError({ storeName: 'pxe', cause })`. No cleanup
+ *     needed (nothing was opened).
+ *   - Wallet store fails to decrypt → closes the already-opened PXE store then throws
+ *     `EmbeddedWalletEncryptionError({ storeName: 'wallet', cause })`.
+ *   - Any non-decrypt error during the wallet open → still closes PXE, then re-throws the original error unwrapped
+ *     (preserves callers' existing untyped error handling for non-encryption faults).
+ *
+ * @param config           - Per-store name/poolDirectory.
+ * @param getEncryptionKey - Returns a fresh 32-byte key per call (the buffer
+ *                           detaches on transfer, so each call must allocate).
+ * @param log              - Logger for both stores.
+ * @param openStore        - Internal test seam. Do not pass in production code.
+ */
+export async function openEncryptedEmbeddedStores(
+  config: OpenEncryptedEmbeddedStoresOptions,
+  getEncryptionKey: () => Promise<Uint8Array>,
+  log: Logger,
+  openStore: OpenSqliteEncryptedStoreFn = defaultOpenStore,
+): Promise<{ pxeStore: AztecSQLiteOPFSStore; walletStore: AztecSQLiteOPFSStore }> {
+  const pxeStore = await openOneStore('pxe', config.pxe, getEncryptionKey, log, openStore);
+  try {
+    const walletStore = await openOneStore('wallet', config.wallet, getEncryptionKey, log, openStore);
+    return { pxeStore, walletStore };
+  } catch (err) {
+    // Cleanup is best-effort — if close() itself throws (e.g. worker already dead), swallow it so the original error
+    // surfaces unobstructed.
+    await pxeStore.close().catch(() => {});
+    throw err;
+  }
+}
+
+async function openOneStore(
+  storeName: EmbeddedStoreName,
+  { name, poolDirectory }: { name: string; poolDirectory?: string },
+  getEncryptionKey: () => Promise<Uint8Array>,
+  log: Logger,
+  openStore: OpenSqliteEncryptedStoreFn,
+): Promise<AztecSQLiteOPFSStore> {
+  const key = await getEncryptionKey();
+  try {
+    return await openStore(log, name, poolDirectory, key);
+  } catch (err) {
+    if (err instanceof SqliteEncryptionError && err.code === 'decrypt_failed') {
+      throw new EmbeddedWalletEncryptionError(storeName, { cause: err });
+    }
+    throw err;
+  }
+}

package/src/embedded/wallet_db.ts

@@ -12,16 +12,15 @@ function accountKey(field: string, address: AztecAddress | string): string {
 }
 
 export class WalletDB {
-  private constructor(
-    private accounts: AztecAsyncMap<string, Buffer>,
-    private aliases: AztecAsyncMap<string, Buffer>,
-    private userLog: LogFn,
-  ) {}
+  private accounts: AztecAsyncMap<string, Buffer>;
+  private aliases: AztecAsyncMap<string, Buffer>;
 
-  static init(store: AztecAsyncKVStore, userLog: LogFn) {
-    const accounts = store.openMap<string, Buffer>('accounts');
-    const aliases = store.openMap<string, Buffer>('aliases');
-    return new WalletDB(accounts, aliases, userLog);
+  constructor(
+    private store: AztecAsyncKVStore,
+    private userLog: LogFn,
+  ) {
+    this.accounts = store.openMap<string, Buffer>('accounts');
+    this.aliases = store.openMap<string, Buffer>('aliases');
   }
 
   async storeAccount(
@@ -131,4 +130,8 @@ export class WalletDB {
       await this.aliases.delete(`accounts:${alias}`);
     }
   }
+
+  async close() {
+    await this.store.close();
+  }
 }