@aztec/wallets 0.0.1-commit.ffe5b04ea → 0.0.1-dev

package/src/embedded/embedded_wallet.ts

@@ -1,33 +1,81 @@
-import { type Account, SignerlessAccount } from '@aztec/aztec.js/account';
-import type { Aliased } from '@aztec/aztec.js/wallet';
-import { AccountManager } from '@aztec/aztec.js/wallet';
+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 { 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 { AccessScopes, PXEConfig, PXECreationOptions } from '@aztec/pxe/client/lazy';
+import type { AztecAsyncKVStore } from '@aztec/kv-store';
+import type { PXEConfig, PXECreationOptions } from '@aztec/pxe/client/lazy';
 import type { PXE } from '@aztec/pxe/server';
 import { AztecAddress } from '@aztec/stdlib/aztec-address';
 import { getContractInstanceFromInstantiationParams } from '@aztec/stdlib/contract';
+import { GasSettings } from '@aztec/stdlib/gas';
 import type { AztecNode } from '@aztec/stdlib/interfaces/client';
 import { deriveSigningKey } from '@aztec/stdlib/keys';
-import { ExecutionPayload, type TxSimulationResult, mergeExecutionPayloads } from '@aztec/stdlib/tx';
-import { BaseWallet, type FeeOptions } from '@aztec/wallet-sdk/base-wallet';
+import {
+  type ContractOverrides,
+  ExecutionPayload,
+  SimulationOverrides,
+  type TxExecutionRequest,
+  TxStatus,
+  collectOffchainEffects,
+  mergeExecutionPayloads,
+} from '@aztec/stdlib/tx';
+import { BaseWallet, type SimulateViaEntrypointOptions } from '@aztec/wallet-sdk/base-wallet';
 
 import type { AccountContractsProvider } from './account-contract-providers/types.js';
 import { type AccountType, WalletDB } from './wallet_db.js';
 
+/** Options for the PXE instance created by the EmbeddedWallet. */
+export type EmbeddedWalletPXEOptions = Partial<PXEConfig> & PXECreationOptions;
+
+/** Splits a unified EmbeddedWalletPXEOptions into PXEConfig overrides and PXECreationOptions. */
+export function splitPxeOptions(pxe?: EmbeddedWalletPXEOptions): {
+  config: Partial<PXEConfig>;
+  creation: PXECreationOptions;
+} {
+  if (!pxe) {
+    return { config: {}, creation: {} };
+  }
+  const { loggers, loggerActorLabel, proverOrOptions, store, simulator, ...config } = pxe;
+  return { config, creation: { loggers, loggerActorLabel, proverOrOptions, store, simulator } };
+}
+
+/** 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;
   /** Use ephemeral (in-memory) stores. Data will not persist across sessions. */
   ephemeral?: boolean;
-  /** Override PXE configuration. */
+  /** 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.
+   */
   pxeConfig?: Partial<PXEConfig>;
-  /** Advanced PXE creation options (custom store, prover, simulator). */
+  /**
+   * Advanced PXE creation options (custom store, prover, simulator).
+   * @deprecated Use `pxe` instead.
+   */
   pxeOptions?: PXECreationOptions;
 };
 
+const DEFAULT_ESTIMATED_GAS_PADDING = 0.1;
+
 export class EmbeddedWallet extends BaseWallet {
+  protected estimatedGasPadding = DEFAULT_ESTIMATED_GAS_PADDING;
+
   constructor(
     pxe: PXE,
     aztecNode: AztecNode,
@@ -39,10 +87,6 @@ export class EmbeddedWallet extends BaseWallet {
   }
 
   protected async getAccountFromAddress(address: AztecAddress): Promise<Account> {
-    if (address.equals(AztecAddress.ZERO)) {
-      return new SignerlessAccount();
-    }
-
     const { secretKey, salt, signingKey, type } = await this.walletDB.retrieveAccount(address);
     const accountManager = await this.createAccountInternal(type, secretKey, salt, signingKey);
     const account = await accountManager.getAccount();
@@ -75,81 +119,172 @@ export class EmbeddedWallet extends BaseWallet {
   }
 
   /**
-   * Simulates calls via a stub account entrypoint, bypassing real account authorization.
-   * This allows kernelless simulation with contract overrides, skipping expensive
-   * private kernel circuit execution.
+   * Overrides the base sendTx to add a pre-simulation step before the actual send. The simulation
+   * estimates actual gas usage and captures call authorization requests to generate
+   * the necessary authwitnesses.
    */
-  protected override async simulateViaEntrypoint(
+  public override async sendTx<W extends InteractionWaitOptions = undefined>(
     executionPayload: ExecutionPayload,
-    from: AztecAddress,
-    feeOptions: FeeOptions,
-    scopes: AccessScopes,
-    _skipTxValidation?: boolean,
-    _skipFeeEnforcement?: boolean,
-  ): Promise<TxSimulationResult> {
-    const { account: fromAccount, instance, artifact } = await this.getFakeAccountDataFor(from);
+    opts: SendOptions<W>,
+  ): Promise<SendReturn<W>> {
+    const feeOptions = await this.completeFeeOptions({
+      from: opts.from,
+      feePayer: executionPayload.feePayer,
+      gasSettings: opts.fee?.gasSettings,
+      forEstimation: true,
+    });
 
-    const feeExecutionPayload = await feeOptions.walletFeePaymentMethod?.getExecutionPayload();
-    const executionOptions: DefaultAccountEntrypointOptions = {
-      txNonce: Fr.random(),
-      cancellable: this.cancellableTransactions,
-      feePaymentMethodOptions: feeOptions.accountFeePaymentMethodOptions,
-    };
-    const finalExecutionPayload = feeExecutionPayload
-      ? mergeExecutionPayloads([feeExecutionPayload, executionPayload])
-      : executionPayload;
-    const chainInfo = await this.getChainInfo();
-    const txRequest = await fromAccount.createTxExecutionRequest(
-      finalExecutionPayload,
-      feeOptions.gasSettings,
-      chainInfo,
-      executionOptions,
-    );
-    return this.pxe.simulateTx(txRequest, {
-      simulatePublic: true,
-      skipFeeEnforcement: true,
+    // Simulate the transaction first to estimate gas and capture required
+    // private authwitnesses based on offchain effects.
+    const simulationResult = await this.simulateViaEntrypoint(executionPayload, {
+      from: opts.from,
+      feeOptions,
+      additionalScopes: opts.additionalScopes,
       skipTxValidation: true,
-      overrides: {
-        contracts: { [from.toString()]: { instance, artifact } },
-      },
-      scopes,
+      sendMessagesAs: opts.sendMessagesAs,
+    });
+
+    const offchainEffects = collectOffchainEffects(simulationResult.privateExecutionResult);
+    const authWitnesses = await Promise.all(
+      offchainEffects.map(async effect => {
+        try {
+          const authRequest = await CallAuthorizationRequest.fromFields(effect.data);
+          return this.createAuthWit(authRequest.onBehalfOf, {
+            consumer: effect.contractAddress,
+            innerHash: authRequest.innerHash,
+          });
+        } catch {
+          return undefined;
+        }
+      }),
+    );
+    for (const authwit of authWitnesses) {
+      if (authwit) {
+        executionPayload.authWitnesses.push(authwit);
+      }
+    }
+    const estimated = getGasLimits(simulationResult, this.estimatedGasPadding);
+    this.log.verbose(
+      `Estimated gas limits for tx: DA=${estimated.gasLimits.daGas} L2=${estimated.gasLimits.l2Gas} teardownDA=${estimated.teardownGasLimits.daGas} teardownL2=${estimated.teardownGasLimits.l2Gas}`,
+    );
+    const gasSettings = GasSettings.from({
+      ...opts.fee?.gasSettings,
+      maxFeesPerGas: feeOptions.gasSettings.maxFeesPerGas,
+      maxPriorityFeesPerGas: feeOptions.gasSettings.maxPriorityFeesPerGas,
+      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;
+    }
+    return super.sendTx(executionPayload, {
+      ...opts,
+      fee: { ...opts.fee, gasSettings },
     });
   }
 
-  private async getFakeAccountDataFor(address: AztecAddress) {
-    // While we have the convention of "Zero address means no auth", and also
-    // we don't have a way to trigger kernelless simulations without overrides,
-    // we need to explicitly handle the zero address case here by
-    // returning the actual multicall contract instead of trying to create a stub account for it.
-    if (!address.equals(AztecAddress.ZERO)) {
+  /**
+   * 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.
+   */
+  protected async buildAccountOverrides(addresses: AztecAddress[]): Promise<ContractOverrides> {
+    const accounts = await this.getAccounts();
+    const contracts: ContractOverrides = {};
+
+    const filtered = accounts.filter(acc => addresses.some(addr => addr.equals(acc.item)));
+
+    for (const account of filtered) {
+      const address = account.item;
+      const { type } = await this.walletDB.retrieveAccount(address);
+      const stubArtifact = await this.accountContracts.getStubAccountContractArtifact(type);
+
       const originalAccount = await this.getAccountFromAddress(address);
-      if (originalAccount instanceof SignerlessAccount) {
-        throw new Error(`Cannot create fake account data for SignerlessAccount at address: ${address}`);
-      }
-      const originalAddress = (originalAccount as Account).getCompleteAddress();
-      const contractInstance = await this.pxe.getContractInstance(originalAddress.address);
+      const completeAddress = originalAccount.getCompleteAddress();
+      const contractInstance = await this.pxe.getContractInstance(completeAddress.address);
       if (!contractInstance) {
-        throw new Error(`No contract instance found for address: ${originalAddress.address}`);
+        throw new Error(
+          `No contract instance found for address: ${completeAddress.address} during account override building. This is a bug!`,
+        );
       }
-      const stubAccount = await this.accountContracts.createStubAccount(originalAddress);
-      const stubArtifact = await this.accountContracts.getStubAccountContractArtifact();
-      const instance = await getContractInstanceFromInstantiationParams(stubArtifact, {
+
+      const stubConstructorArgs = type === 'schnorr' ? [Fr.ZERO, Fr.ZERO] : [Buffer.alloc(32), Buffer.alloc(32)];
+      const stubInstance = await getContractInstanceFromInstantiationParams(stubArtifact, {
         salt: Fr.random(),
+        constructorArgs: stubConstructorArgs,
       });
-      return {
-        account: stubAccount,
-        instance,
+
+      contracts[address.toString()] = {
+        instance: stubInstance,
         artifact: stubArtifact,
       };
+    }
+
+    return contracts;
+  }
+
+  /**
+   * Simulates calls via a stub account entrypoint, bypassing real account authorization.
+   * This allows kernelless simulation with contract overrides, skipping expensive
+   * private kernel circuit execution.
+   */
+  protected override async simulateViaEntrypoint(
+    executionPayload: ExecutionPayload,
+    opts: SimulateViaEntrypointOptions,
+  ): Promise<TxSimulationResultWithAppOffset> {
+    const { from, feeOptions, additionalScopes, skipTxValidation, skipFeeEnforcement, sendMessagesAs } = opts;
+    const scopes = this.scopesFrom(from, additionalScopes);
+
+    const feeExecutionPayload = await feeOptions.walletFeePaymentMethod?.getExecutionPayload();
+    const finalExecutionPayload = feeExecutionPayload
+      ? mergeExecutionPayloads([feeExecutionPayload, executionPayload])
+      : executionPayload;
+    const chainInfo = await this.getChainInfo();
+
+    const accountOverrides = await this.buildAccountOverrides(scopes);
+    const overrides = new SimulationOverrides(accountOverrides);
+
+    let txRequest: TxExecutionRequest;
+    if (from === NO_FROM) {
+      const entrypoint = new DefaultEntrypoint();
+      txRequest = await entrypoint.createTxExecutionRequest(finalExecutionPayload, feeOptions.gasSettings, chainInfo);
     } else {
-      const { instance, artifact } = await this.accountContracts.getMulticallContract();
-      const account = new SignerlessAccount();
-      return {
-        instance,
-        account,
-        artifact,
+      const { type } = await this.walletDB.retrieveAccount(from);
+      const originalAccount = await this.getAccountFromAddress(from);
+      const completeAddress = originalAccount.getCompleteAddress();
+      const account = await this.accountContracts.createStubAccount(completeAddress, type);
+      const executionOptions: DefaultAccountEntrypointOptions = {
+        txNonce: Fr.random(),
+        cancellable: this.cancellableTransactions,
+        // If from is an address, feeOptions include the way the account contract should handle the fee payment
+        feePaymentMethodOptions: feeOptions.accountFeePaymentMethodOptions!,
       };
+      txRequest = await account.createTxExecutionRequest(
+        finalExecutionPayload,
+        feeOptions.gasSettings,
+        chainInfo,
+        executionOptions,
+      );
     }
+
+    const result = await this.pxe.simulateTx(txRequest, {
+      simulatePublic: true,
+      skipFeeEnforcement,
+      skipTxValidation,
+      overrides,
+      scopes,
+      senderForTags: this.senderForTagsFrom(from, sendMessagesAs),
+    });
+    const appCallOffset = await this.computeAppCallOffset(from, feeOptions);
+    return TxSimulationResultWithAppOffset.fromResultAndOffset(result, appCallOffset);
   }
 
   protected async createAccountInternal(
@@ -221,7 +356,12 @@ export class EmbeddedWallet extends BaseWallet {
     this.minFeePadding = value ?? 0.5;
   }
 
-  stop() {
-    return this.pxe.stop();
+  setEstimatedGasPadding(value?: number) {
+    this.estimatedGasPadding = value ?? DEFAULT_ESTIMATED_GAS_PADDING;
+  }
+
+  async stop(): Promise<void> {
+    await this.pxe.stop();
+    await this.walletDB.close();
   }
 }

package/src/embedded/entrypoints/browser.ts

@@ -6,7 +6,7 @@ import { type PXEConfig, getPXEConfig } from '@aztec/pxe/config';
 
 import { LazyAccountContractsProvider } from '../account-contract-providers/lazy.js';
 import type { AccountContractsProvider } from '../account-contract-providers/types.js';
-import { EmbeddedWallet, type EmbeddedWalletOptions } from '../embedded_wallet.js';
+import { EmbeddedWallet, type EmbeddedWalletOptions, splitPxeOptions } from '../embedded_wallet.js';
 import { WalletDB } from '../wallet_db.js';
 
 export class BrowserEmbeddedWallet extends EmbeddedWallet {
@@ -26,10 +26,15 @@ export class BrowserEmbeddedWallet extends EmbeddedWallet {
     const aztecNode = typeof nodeOrUrl === 'string' ? createAztecNodeClient(nodeOrUrl) : nodeOrUrl;
     const l1Contracts = await aztecNode.getL1ContractAddresses();
 
+    // Support both the new unified `pxe` option and the deprecated `pxeConfig`/`pxeOptions`.
+    const { config: pxeConfigFromPxe, creation: pxeCreationFromPxe } = splitPxeOptions(options.pxe);
+    const mergedConfigOverrides = { ...options.pxeConfig, ...pxeConfigFromPxe };
+    const mergedCreationOverrides: PXECreationOptions = { ...options.pxeOptions, ...pxeCreationFromPxe };
+
     const pxeConfig: PXEConfig = Object.assign(getPXEConfig(), {
-      proverEnabled: options.pxeConfig?.proverEnabled ?? false,
+      proverEnabled: mergedConfigOverrides.proverEnabled ?? false,
       dataDirectory: `pxe_data_${l1Contracts.rollupAddress}`,
-      ...options.pxeConfig,
+      ...mergedConfigOverrides,
     });
 
     if (options.ephemeral) {
@@ -37,36 +42,45 @@ export class BrowserEmbeddedWallet extends EmbeddedWallet {
     }
 
     const pxeOptions: PXECreationOptions = {
-      ...options.pxeOptions,
+      ...mergedCreationOverrides,
       loggers: {
         store: rootLogger.createChild('pxe:data'),
         pxe: rootLogger.createChild('pxe:service'),
         prover: rootLogger.createChild('pxe:prover'),
-        ...options.pxeOptions?.loggers,
+        ...mergedCreationOverrides.loggers,
       },
     };
 
     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,
+              l1Contracts,
+            },
+            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;
   }
 }
 
 export { BrowserEmbeddedWallet as EmbeddedWallet };
-export type { EmbeddedWalletOptions } from '../embedded_wallet.js';
+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

@@ -7,7 +7,7 @@ import type { AztecNode } from '@aztec/stdlib/interfaces/client';
 
 import { BundleAccountContractsProvider } from '../account-contract-providers/bundle.js';
 import type { AccountContractsProvider } from '../account-contract-providers/types.js';
-import { EmbeddedWallet, type EmbeddedWalletOptions } from '../embedded_wallet.js';
+import { EmbeddedWallet, type EmbeddedWalletOptions, splitPxeOptions } from '../embedded_wallet.js';
 import { WalletDB } from '../wallet_db.js';
 
 export class NodeEmbeddedWallet extends EmbeddedWallet {
@@ -27,10 +27,15 @@ export class NodeEmbeddedWallet extends EmbeddedWallet {
     const aztecNode = typeof nodeOrUrl === 'string' ? createAztecNodeClient(nodeOrUrl) : nodeOrUrl;
     const l1Contracts = await aztecNode.getL1ContractAddresses();
 
+    // Support both the new unified `pxe` option and the deprecated `pxeConfig`/`pxeOptions`.
+    const { config: pxeConfigFromPxe, creation: pxeCreationFromPxe } = splitPxeOptions(options.pxe);
+    const mergedConfigOverrides = { ...options.pxeConfig, ...pxeConfigFromPxe };
+    const mergedCreationOverrides: PXECreationOptions = { ...options.pxeOptions, ...pxeCreationFromPxe };
+
     const pxeConfig: PXEConfig = Object.assign(getPXEConfig(), {
-      proverEnabled: options.pxeConfig?.proverEnabled ?? false,
+      proverEnabled: mergedConfigOverrides.proverEnabled ?? false,
       dataDirectory: `pxe_data_${l1Contracts.rollupAddress}`,
-      ...options.pxeConfig,
+      ...mergedConfigOverrides,
     });
 
     if (options.ephemeral) {
@@ -38,42 +43,44 @@ export class NodeEmbeddedWallet extends EmbeddedWallet {
     }
 
     const pxeOptions: PXECreationOptions = {
-      ...options.pxeOptions,
+      ...mergedCreationOverrides,
       loggers: {
         store: rootLogger.createChild('pxe:data'),
         pxe: rootLogger.createChild('pxe:service'),
         prover: rootLogger.createChild('pxe:prover'),
-        ...options.pxeOptions?.loggers,
+        ...mergedCreationOverrides.loggers,
       },
     };
 
     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,
+              l1Contracts,
+            },
+            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;
   }
 }
 
 export { NodeEmbeddedWallet as EmbeddedWallet };
-export type { EmbeddedWalletOptions } from '../embedded_wallet.js';
+export type { EmbeddedWalletOptions, EmbeddedWalletPXEOptions } from '../embedded_wallet.js';
 export { WalletDB } from '../wallet_db.js';
 export type { AccountType } from '../wallet_db.js';

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;
+  }
+}