fix(pxe): support custom PrivateKernelProver and unify EmbeddedWalletOptions (backport #22348)

docs/docs-developers/docs/resources/migration_notes.md

@@ -91,6 +91,22 @@ FPCs that use only Fee Juice still work on all networks, since FeeJuice is a pro
 
 Similarly, the `fpc-public` and `fpc-private` CLI wallet payment methods use the reference Token-based FPC and will not work on public networks. Use `fee_juice` for direct Fee Juice payment, or `fpc-sponsored` on devnet and local network.
 
+### [aztec.js] `EmbeddedWalletOptions` now uses a unified `pxe` field
+
+The `pxeConfig` and `pxeOptions` fields on `EmbeddedWalletOptions` have been deprecated in favor of a single `pxe` field that accepts both PXE configuration and dependency overrides (custom prover, store, simulator):
+
+```diff
+const wallet = await EmbeddedWallet.create(nodeUrl, {
+-  pxeConfig: { proverEnabled: true },
+-  pxeOptions: { proverOrOptions: myCustomProver },
++  pxe: {
++    proverEnabled: true,
++    proverOrOptions: myCustomProver,
++  },
+});
+```
+
+The old fields still work but will be removed in a future release.
 
 ### [Aztec.nr] Domain-separated tags on log emission
 

yarn-project/pxe/src/entrypoints/client/bundle/utils.ts

@@ -1,4 +1,3 @@
-import { BBPrivateKernelProver } from '@aztec/bb-prover/client';
 import { BBBundlePrivateKernelProver } from '@aztec/bb-prover/client/bundle';
 import { createLogger } from '@aztec/foundation/log';
 import { createStore } from '@aztec/kv-store/indexeddb';
@@ -9,7 +8,7 @@ import type { AztecNode } from '@aztec/stdlib/interfaces/client';
 import type { PXEConfig } from '../../../config/index.js';
 import { PXE } from '../../../pxe.js';
 import { PXE_DATA_SCHEMA_VERSION } from '../../../storage/metadata.js';
-import type { PXECreationOptions } from '../../pxe_creation_options.js';
+import { type PXECreationOptions, isPrivateKernelProver } from '../../pxe_creation_options.js';
 
 /**
  * Create and start an PXE instance with the given AztecNode.
@@ -44,7 +43,7 @@ export async function createPXE(
   const proverLogger = loggers.prover ?? createLogger('pxe:bb:wasm:bundle', { actor });
 
   let prover;
-  if (options.proverOrOptions instanceof BBPrivateKernelProver) {
+  if (isPrivateKernelProver(options.proverOrOptions)) {
     prover = options.proverOrOptions;
   } else {
     prover = new BBBundlePrivateKernelProver(simulator, { ...options.proverOrOptions, logger: proverLogger });

yarn-project/pxe/src/entrypoints/client/lazy/utils.ts

@@ -1,4 +1,3 @@
-import { BBPrivateKernelProver } from '@aztec/bb-prover/client';
 import { BBLazyPrivateKernelProver } from '@aztec/bb-prover/client/lazy';
 import { createLogger } from '@aztec/foundation/log';
 import { createStore } from '@aztec/kv-store/indexeddb';
@@ -9,7 +8,7 @@ import type { AztecNode } from '@aztec/stdlib/interfaces/client';
 import type { PXEConfig } from '../../../config/index.js';
 import { PXE } from '../../../pxe.js';
 import { PXE_DATA_SCHEMA_VERSION } from '../../../storage/metadata.js';
-import type { PXECreationOptions } from '../../pxe_creation_options.js';
+import { type PXECreationOptions, isPrivateKernelProver } from '../../pxe_creation_options.js';
 
 /**
  * Create and start an PXE instance with the given AztecNode.
@@ -44,7 +43,7 @@ export async function createPXE(
   const proverLogger = loggers.prover ?? createLogger('pxe:bb:wasm:bundle', { actor });
 
   let prover;
-  if (options.proverOrOptions instanceof BBPrivateKernelProver) {
+  if (isPrivateKernelProver(options.proverOrOptions)) {
     prover = options.proverOrOptions;
   } else {
     prover = new BBLazyPrivateKernelProver(simulator, { ...options.proverOrOptions, logger: proverLogger });

yarn-project/pxe/src/entrypoints/pxe_creation_options.ts

@@ -12,3 +12,10 @@ export type PXECreationOptions = {
   store?: AztecAsyncKVStore;
   simulator?: CircuitSimulator;
 };
+
+/** Checks if the given value implements the PrivateKernelProver interface via duck-typing. */
+export function isPrivateKernelProver(value: unknown): value is PrivateKernelProver {
+  return (
+    typeof value === 'object' && value !== null && typeof (value as PrivateKernelProver).createChonkProof === 'function'
+  );
+}

yarn-project/pxe/src/entrypoints/server/utils.ts

@@ -1,4 +1,3 @@
-import { BBPrivateKernelProver } from '@aztec/bb-prover/client';
 import { BBBundlePrivateKernelProver } from '@aztec/bb-prover/client/bundle';
 import { createLogger } from '@aztec/foundation/log';
 import { createStore } from '@aztec/kv-store/lmdb-v2';
@@ -10,7 +9,7 @@ import type { AztecNode } from '@aztec/stdlib/interfaces/client';
 import type { PXEConfig } from '../../config/index.js';
 import { PXE } from '../../pxe.js';
 import { PXE_DATA_SCHEMA_VERSION } from '../../storage/index.js';
-import type { PXECreationOptions } from '../pxe_creation_options.js';
+import { type PXECreationOptions, isPrivateKernelProver } from '../pxe_creation_options.js';
 
 type PXEConfigWithoutDefaults = Omit<PXEConfig, 'l1Contracts' | 'l1ChainId' | 'l2BlockBatchSize' | 'rollupVersion'>;
 
@@ -49,7 +48,7 @@ export async function createPXE(
   const proverLogger = loggers.prover ?? createLogger('pxe:bb:native', { actor });
 
   let prover;
-  if (options.proverOrOptions instanceof BBPrivateKernelProver) {
+  if (isPrivateKernelProver(options.proverOrOptions)) {
     prover = options.proverOrOptions;
   } else {
     prover = new BBBundlePrivateKernelProver(simulator, { ...options.proverOrOptions, logger: proverLogger });

yarn-project/wallets/src/embedded/embedded_wallet.ts

@@ -29,14 +29,37 @@ import { BaseWallet, type SimulateViaEntrypointOptions } from '@aztec/wallet-sdk
 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 } };
+}
+
 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;
+  /**
+   * 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;
 };
 

yarn-project/wallets/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,23 +26,28 @@ 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) {
       delete pxeConfig.dataDirectory;
     }
 
     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,
       },
     };
 
@@ -67,6 +72,6 @@ export class BrowserEmbeddedWallet extends EmbeddedWallet {
 }
 
 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';

yarn-project/wallets/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,23 +27,28 @@ 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) {
       delete pxeConfig.dataDirectory;
     }
 
     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,
       },
     };
 
@@ -74,6 +79,6 @@ export class NodeEmbeddedWallet extends EmbeddedWallet {
 }
 
 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';