@aztec/aztec.js 3.0.0-canary.a9708bd → 3.0.0-devnet.2

package/src/fee/public_fee_payment_method.ts

@@ -1,12 +1,13 @@
-import type { FeePaymentMethod } from '@aztec/entrypoints/interfaces';
 import { ExecutionPayload } from '@aztec/entrypoints/payload';
 import { Fr } from '@aztec/foundation/fields';
-import { FunctionSelector, FunctionType } from '@aztec/stdlib/abi';
+import { type FunctionAbi, FunctionSelector, FunctionType, decodeFromAbi } from '@aztec/stdlib/abi';
 import { AztecAddress } from '@aztec/stdlib/aztec-address';
 import { GasSettings } from '@aztec/stdlib/gas';
 
-import type { AccountWallet } from '../wallet/account_wallet.js';
-import { simulateWithoutSignature } from './utils.js';
+import { ContractFunctionInteraction } from '../contract/contract_function_interaction.js';
+import { SetPublicAuthwitContractInteraction } from '../utils/authwit.js';
+import type { Wallet } from '../wallet/wallet.js';
+import type { FeePaymentMethod } from './fee_payment_method.js';
 
 /**
  * Holds information about how the fee for a transaction is to be paid.
@@ -22,45 +23,58 @@ export class PublicFeePaymentMethod implements FeePaymentMethod {
     /**
      * An auth witness provider to authorize fee payments
      */
-    protected wallet: AccountWallet,
+    protected sender: AztecAddress,
+    /**
+     * A wallet to perform the simulation to get the accepted asset
+     */
+    protected wallet: Wallet,
+    /**
+     * Gas settings used to compute the maximum fee the user is willing to pay
+     */
+    protected gasSettings: GasSettings,
   ) {}
 
   /**
    * The asset used to pay the fee.
    * @returns The asset used to pay the fee.
    */
-  getAsset(): Promise<AztecAddress> {
+  async getAsset(): Promise<AztecAddress> {
     if (!this.assetPromise) {
-      // We use the utility method to avoid a signature because this function could be triggered
-      // before the associated account is deployed.
-      this.assetPromise = simulateWithoutSignature(
-        this.wallet,
-        this.paymentContract,
-        {
-          name: 'get_accepted_asset',
-          functionType: FunctionType.PRIVATE,
-          isInternal: false,
-          isStatic: false,
-          parameters: [],
-          returnTypes: [
-            {
-              kind: 'struct',
-              path: 'authwit::aztec::protocol_types::address::aztec_address::AztecAddress',
-              fields: [
-                {
-                  name: 'inner',
-                  type: {
-                    kind: 'field',
-                  },
+      const abi = {
+        name: 'get_accepted_asset',
+        functionType: FunctionType.PRIVATE,
+        isInternal: false,
+        isStatic: false,
+        parameters: [],
+        returnTypes: [
+          {
+            kind: 'struct',
+            path: 'authwit::aztec::protocol_types::address::aztec_address::AztecAddress',
+            fields: [
+              {
+                name: 'inner',
+                type: {
+                  kind: 'field',
                 },
-              ],
-            },
-          ],
-          errorTypes: {},
-          isInitializer: false,
-        },
-        [],
-      ) as Promise<AztecAddress>;
+              },
+            ],
+          },
+        ],
+        errorTypes: {},
+        isInitializer: false,
+      } as FunctionAbi;
+      const interaction = new ContractFunctionInteraction(this.wallet, this.paymentContract, abi, []);
+
+      const executionPayload = await interaction.request();
+      this.assetPromise = this.wallet
+        .simulateTx(executionPayload, {
+          from: AztecAddress.ZERO,
+          skipFeeEnforcement: true,
+        })
+        .then(simulationResult => {
+          const rawReturnValues = simulationResult.getPrivateReturnValues().nested[0].values;
+          return decodeFromAbi(abi.returnTypes, rawReturnValues!);
+        }) as Promise<AztecAddress>;
     }
     return this.assetPromise!;
   }
@@ -74,23 +88,28 @@ export class PublicFeePaymentMethod implements FeePaymentMethod {
    * @param gasSettings - The gas settings.
    * @returns An execution payload that contains the required function calls.
    */
-  async getExecutionPayload(gasSettings: GasSettings): Promise<ExecutionPayload> {
+  async getExecutionPayload(): Promise<ExecutionPayload> {
     const txNonce = Fr.random();
-    const maxFee = gasSettings.getFeeLimit();
+    const maxFee = this.gasSettings.getFeeLimit();
 
-    const setPublicAuthWitInteraction = await this.wallet.setPublicAuthWit(
-      {
-        caller: this.paymentContract,
-        action: {
-          name: 'transfer_in_public',
-          args: [this.wallet.getAddress().toField(), this.paymentContract.toField(), maxFee, txNonce],
-          selector: await FunctionSelector.fromSignature('transfer_in_public((Field),(Field),u128,Field)'),
-          type: FunctionType.PUBLIC,
-          isStatic: false,
-          to: await this.getAsset(),
-          returnTypes: [],
-        },
+    const intent = {
+      caller: this.paymentContract,
+      call: {
+        name: 'transfer_in_public',
+        args: [this.sender.toField(), this.paymentContract.toField(), maxFee, txNonce],
+        selector: await FunctionSelector.fromSignature('transfer_in_public((Field),(Field),u128,Field)'),
+        type: FunctionType.PUBLIC,
+        isStatic: false,
+        hideMsgSender: false /** The target function performs an authwit check, so msg_sender is needed */,
+        to: await this.getAsset(),
+        returnTypes: [],
       },
+    };
+
+    const setPublicAuthWitInteraction = await SetPublicAuthwitContractInteraction.create(
+      this.wallet,
+      this.sender,
+      intent,
       true,
     );
 
@@ -102,6 +121,7 @@ export class PublicFeePaymentMethod implements FeePaymentMethod {
           to: this.paymentContract,
           selector: await FunctionSelector.fromSignature('fee_entrypoint_public(u128,Field)'),
           type: FunctionType.PRIVATE,
+          hideMsgSender: false,
           isStatic: false,
           args: [maxFee, txNonce],
           returnTypes: [],
@@ -111,4 +131,8 @@ export class PublicFeePaymentMethod implements FeePaymentMethod {
       [],
     );
   }
+
+  getGasSettings(): GasSettings | undefined {
+    return this.gasSettings;
+  }
 }

package/src/fee/sponsored_fee_payment.ts

@@ -2,6 +2,7 @@ import type { FeePaymentMethod } from '@aztec/aztec.js/fee';
 import { ExecutionPayload } from '@aztec/entrypoints/payload';
 import { FunctionSelector, FunctionType } from '@aztec/stdlib/abi';
 import { AztecAddress } from '@aztec/stdlib/aztec-address';
+import type { GasSettings } from '@aztec/stdlib/gas';
 
 /**
  * A fee payment method that uses a contract that blindly sponsors transactions.
@@ -26,6 +27,7 @@ export class SponsoredFeePaymentMethod implements FeePaymentMethod {
           to: this.paymentContract,
           selector: await FunctionSelector.fromSignature('sponsor_unconditionally()'),
           type: FunctionType.PRIVATE,
+          hideMsgSender: false,
           isStatic: false,
           args: [],
           returnTypes: [],
@@ -35,4 +37,8 @@ export class SponsoredFeePaymentMethod implements FeePaymentMethod {
       [],
     );
   }
+
+  getGasSettings(): GasSettings | undefined {
+    return;
+  }
 }

package/src/utils/authwit.ts

@@ -1,36 +1,57 @@
+import type { ChainInfo } from '@aztec/entrypoints/interfaces';
 import { Fr } from '@aztec/foundation/fields';
-import type { FunctionCall } from '@aztec/stdlib/abi';
-import { computeInnerAuthWitHash, computeOuterAuthWitHash } from '@aztec/stdlib/auth-witness';
+import { ProtocolContractAddress } from '@aztec/protocol-contracts';
+import { type ABIParameterVisibility, type FunctionAbi, type FunctionCall, FunctionType } from '@aztec/stdlib/abi';
+import { AuthWitness, computeInnerAuthWitHash, computeOuterAuthWitHash } from '@aztec/stdlib/auth-witness';
 import type { AztecAddress } from '@aztec/stdlib/aztec-address';
 import { computeVarArgsHash } from '@aztec/stdlib/hash';
+import type { TxProfileResult } from '@aztec/stdlib/tx';
 
 import { ContractFunctionInteraction } from '../contract/contract_function_interaction.js';
-
-/** Metadata for the intent */
-export type IntentMetadata = {
-  /** The chain id to approve */
-  chainId: Fr;
-  /** The version to approve  */
-  version: Fr;
-};
+import type {
+  ProfileInteractionOptions,
+  SendInteractionOptions,
+  SimulateInteractionOptions,
+  SimulationReturn,
+} from '../contract/interaction_options.js';
+import type { SentTx } from '../contract/sent_tx.js';
+import type { Wallet } from '../wallet/index.js';
 
 /** Intent with an inner hash */
 export type IntentInnerHash = {
   /** The consumer   */
   consumer: AztecAddress;
   /** The action to approve */
-  innerHash: Buffer | Fr;
+  innerHash: Buffer<ArrayBuffer> | Fr;
+};
+
+/** Intent with a call */
+export type CallIntent = {
+  /** The caller to approve  */
+  caller: AztecAddress;
+  /** The call to approve */
+  call: FunctionCall;
 };
 
-/** Intent with an action */
-export type IntentAction = {
+/** Intent with a ContractFunctionInteraction */
+export type ContractFunctionInteractionCallIntent = {
   /** The caller to approve  */
   caller: AztecAddress;
   /** The action to approve */
-  action: ContractFunctionInteraction | FunctionCall;
+  action: ContractFunctionInteraction;
 };
 
-// docs:start:authwit_computeAuthWitMessageHash
+/** Identifies ContractFunctionInteractionCallIntents */
+function isContractFunctionIntractionCallIntent(
+  messageHashOrIntent: Fr | Buffer | IntentInnerHash | CallIntent | ContractFunctionInteractionCallIntent,
+): messageHashOrIntent is ContractFunctionInteractionCallIntent {
+  return (
+    'caller' in messageHashOrIntent &&
+    'action' in messageHashOrIntent &&
+    messageHashOrIntent.action instanceof ContractFunctionInteraction
+  );
+}
+
 /**
  * Compute an authentication witness message hash from an intent and metadata
  *
@@ -41,7 +62,7 @@ export type IntentAction = {
  * and use it for the authentication check.
  * Therefore, any allowed `innerHash` will therefore also have information around where it can be spent (version, chainId) and who can spend it (consumer).
  *
- * If using the `IntentAction`, the caller is the address that is making the call, for a token approval from Alice to Bob, this would be Bob.
+ * If using the `CallIntent`, the caller is the address that is making the call, for a token approval from Alice to Bob, this would be Bob.
  * The action is then used along with the `caller` to compute the `innerHash` and the consumer.
  *
  *
@@ -52,49 +73,240 @@ export type IntentAction = {
  * @param metadata - The metadata for the intent (chainId, version)
  * @returns The message hash for the action
  */
-export const computeAuthWitMessageHash = async (intent: IntentInnerHash | IntentAction, metadata: IntentMetadata) => {
+export const computeAuthWitMessageHash = async (
+  intent: IntentInnerHash | CallIntent | ContractFunctionInteractionCallIntent,
+  metadata: ChainInfo,
+) => {
   const chainId = metadata.chainId;
   const version = metadata.version;
 
   if ('caller' in intent) {
-    const fnCall =
-      intent.action instanceof ContractFunctionInteraction ? (await intent.action.request()).calls[0] : intent.action;
+    const call = isContractFunctionIntractionCallIntent(intent) ? await intent.action.getFunctionCall() : intent.call;
     return computeOuterAuthWitHash(
-      fnCall.to,
+      call.to,
       chainId,
       version,
-      await computeInnerAuthWitHashFromFunctionCall(intent.caller, fnCall),
+      await computeInnerAuthWitHashFromAction(intent.caller, call),
     );
   } else {
     const inner = Buffer.isBuffer(intent.innerHash) ? Fr.fromBuffer(intent.innerHash) : intent.innerHash;
     return computeOuterAuthWitHash(intent.consumer, chainId, version, inner);
   }
 };
-// docs:end:authwit_computeAuthWitMessageHash
 
 /**
- * Computes the inner authwitness hash for a function call, for it to later be combined with the metadata
- * required for the outer hash and eventually the full AuthWitness.
- * @param caller - Who is going to be calling the function
- * @param fnCall - The function call to compute the inner hash from
- * @returns The inner hash for the function call
- **/
-export const computeInnerAuthWitHashFromFunctionCall = async (caller: AztecAddress, fnCall: FunctionCall) => {
-  return computeInnerAuthWitHash([caller.toField(), fnCall.selector.toField(), await computeVarArgsHash(fnCall.args)]);
-};
+ * Compute an authentication witness message hash from an intent and metadata. This is just
+ * a wrapper around computeAuthwitMessageHash that allows receiving an already computed messageHash as input
+ * @param messageHashOrIntent - The precomputed messageHash or intent to approve (consumer and innerHash or caller and call/action)
+ * @param metadata - The metadata for the intent (chainId, version)
+ * @returns The message hash for the intent
+ */
+export async function getMessageHashFromIntent(
+  messageHashOrIntent: Fr | Buffer | IntentInnerHash | CallIntent | ContractFunctionInteractionCallIntent,
+  chainInfo: ChainInfo,
+) {
+  let messageHash: Fr;
+  const { chainId, version } = chainInfo;
+  if (Buffer.isBuffer(messageHashOrIntent)) {
+    messageHash = Fr.fromBuffer(messageHashOrIntent);
+  } else if (messageHashOrIntent instanceof Fr) {
+    messageHash = messageHashOrIntent;
+  } else {
+    messageHash = await computeAuthWitMessageHash(messageHashOrIntent, { chainId, version });
+  }
+  return messageHash;
+}
 
 /**
- * Computes the inner authwitness hash for an action, that can either be a ContractFunctionInteraction
- * or an isolated FunctionCall. Since the former is just a wrapper around the latter, we can just extract
- * the first (and only) call from the ContractFunctionInteraction and use it to compute the inner hash.
- * @param caller - Who is going to be performing the action
- * @param action - The ContractFunctionInteraction or FunctionCall to compute the inner hash for
+ * Computes the inner authwitness hash for either a function call or an action, for it to later be combined with the metadata
+ * required for the outer hash and eventually the full AuthWitness.
+ * @param caller - Who is going to be calling the function
+ * @param action - The action to compute the inner hash from
  * @returns The inner hash for the action
  **/
 export const computeInnerAuthWitHashFromAction = async (
   caller: AztecAddress,
   action: FunctionCall | ContractFunctionInteraction,
 ) => {
-  action = action instanceof ContractFunctionInteraction ? (await action.request()).calls[0] : action;
-  return computeInnerAuthWitHash([caller.toField(), action.selector.toField(), await computeVarArgsHash(action.args)]);
+  const call = action instanceof ContractFunctionInteraction ? await action.getFunctionCall() : action;
+  return computeInnerAuthWitHash([caller.toField(), call.selector.toField(), await computeVarArgsHash(call.args)]);
 };
+
+/**
+ * Lookup the validity of an authwit in private and public contexts.
+ *
+ * Uses the chain id and version of the wallet.
+ *
+ * @param wallet - The wallet use to simulate and read the public data
+ * @param onBehalfOf - The address of the "approver"
+ * @param intent - The consumer and inner hash or the caller and action to lookup
+ * @param witness - The computed authentication witness to check
+ * @returns - A struct containing the validity of the authwit in private and public contexts.
+ */
+export async function lookupValidity(
+  wallet: Wallet,
+  onBehalfOf: AztecAddress,
+  intent: IntentInnerHash | CallIntent | ContractFunctionInteractionCallIntent,
+  witness: AuthWitness,
+): Promise<{
+  /** boolean flag indicating if the authwit is valid in private context */
+  isValidInPrivate: boolean;
+  /** boolean flag indicating if the authwit is valid in public context */
+  isValidInPublic: boolean;
+}> {
+  let innerHash, consumer;
+  if ('caller' in intent) {
+    const call = isContractFunctionIntractionCallIntent(intent) ? await intent.action.getFunctionCall() : intent.call;
+    innerHash = await computeInnerAuthWitHashFromAction(intent.caller, call);
+    consumer = call.to;
+  } else if (Buffer.isBuffer(intent.innerHash)) {
+    innerHash = Fr.fromBuffer(intent.innerHash);
+    consumer = intent.consumer;
+  } else {
+    ({ innerHash, consumer } = intent);
+  }
+  const chainInfo = await wallet.getChainInfo();
+  const messageHash = await getMessageHashFromIntent(intent, chainInfo);
+  const results = { isValidInPrivate: false, isValidInPublic: false };
+
+  // Check private
+  const lookupValidityAbi = {
+    name: 'lookup_validity',
+    isInitializer: false,
+    functionType: FunctionType.UTILITY,
+    isInternal: false,
+    isStatic: false,
+    parameters: [{ name: 'message_hash', type: { kind: 'field' }, visibility: 'private' as ABIParameterVisibility }],
+    returnTypes: [{ kind: 'boolean' }],
+    errorTypes: {},
+  } as FunctionAbi;
+  try {
+    results.isValidInPrivate = (await new ContractFunctionInteraction(wallet, onBehalfOf, lookupValidityAbi, [
+      consumer,
+      innerHash,
+    ]).simulate({ from: onBehalfOf, authWitnesses: [witness] })) as boolean;
+    // TODO: Narrow down the error to make sure simulation failed due to an invalid authwit
+    // eslint-disable-next-line no-empty
+  } catch {}
+
+  // check public
+  const isConsumableAbi = {
+    name: 'utility_is_consumable',
+    isInitializer: false,
+    functionType: FunctionType.UTILITY,
+    isInternal: false,
+    isStatic: false,
+    parameters: [
+      {
+        name: 'address',
+        type: {
+          fields: [{ name: 'inner', type: { kind: 'field' } }],
+          kind: 'struct',
+          path: 'authwit::aztec::protocol_types::address::aztec_address::AztecAddress',
+        },
+        visibility: 'private' as ABIParameterVisibility,
+      },
+      { name: 'message_hash', type: { kind: 'field' }, visibility: 'private' as ABIParameterVisibility },
+    ],
+    returnTypes: [{ kind: 'boolean' }],
+    errorTypes: {},
+  } as FunctionAbi;
+  results.isValidInPublic = (await new ContractFunctionInteraction(
+    wallet,
+    ProtocolContractAddress.AuthRegistry,
+    isConsumableAbi,
+    [onBehalfOf, messageHash],
+  ).simulate({ from: onBehalfOf })) as boolean;
+
+  return results;
+}
+
+/**
+ * Convenience class designed to wrap the very common interaction of setting a public authwit in the AuthRegistry contract
+ */
+export class SetPublicAuthwitContractInteraction extends ContractFunctionInteraction {
+  private constructor(
+    wallet: Wallet,
+    private from: AztecAddress,
+    messageHash: Fr,
+    authorized: boolean,
+  ) {
+    super(wallet, ProtocolContractAddress.AuthRegistry, SetPublicAuthwitContractInteraction.getSetAuthorizedAbi(), [
+      messageHash,
+      authorized,
+    ]);
+  }
+
+  static async create(
+    wallet: Wallet,
+    from: AztecAddress,
+    messageHashOrIntent: Fr | Buffer | IntentInnerHash | CallIntent | ContractFunctionInteractionCallIntent,
+    authorized: boolean,
+  ) {
+    const chainInfo = await wallet.getChainInfo();
+    const messageHash = await getMessageHashFromIntent(messageHashOrIntent, chainInfo);
+    return new SetPublicAuthwitContractInteraction(wallet, from, messageHash, authorized);
+  }
+
+  /**
+   * Overrides the simulate method, adding the sender of the authwit (authorizer) as from
+   * and preventing misuse
+   * @param options - An optional object containing additional configuration for the transaction.
+   * @returns The result of the transaction as returned by the contract function.
+   */
+  public override simulate<T extends SimulateInteractionOptions>(
+    options: Omit<T, 'from'>,
+  ): Promise<SimulationReturn<T['includeMetadata']>>;
+  // eslint-disable-next-line jsdoc/require-jsdoc
+  public override simulate(
+    options: Omit<SimulateInteractionOptions, 'from'> = {},
+  ): Promise<SimulationReturn<typeof options.includeMetadata>> {
+    return super.simulate({ ...options, from: this.from });
+  }
+
+  /**
+   * Overrides the profile method, adding the sender of the authwit (authorizer) as from
+   * and preventing misuse
+   * @param options - Same options as `simulate`, plus profiling method
+   * @returns An object containing the function return value and profile result.
+   */
+  public override profile(
+    options: Omit<ProfileInteractionOptions, 'from'> = { profileMode: 'gates' },
+  ): Promise<TxProfileResult> {
+    return super.profile({ ...options, from: this.from });
+  }
+
+  /**
+   * Overrides the send method, adding the sender of the authwit (authorizer) as from
+   * and preventing misuse
+   * @param options - An optional object containing 'fee' options information
+   * @returns A SentTx instance for tracking the transaction status and information.
+   */
+  public override send(options: Omit<SendInteractionOptions, 'from'> = {}): SentTx {
+    return super.send({ ...options, from: this.from });
+  }
+
+  private static getSetAuthorizedAbi(): FunctionAbi {
+    return {
+      name: 'set_authorized',
+      isInitializer: false,
+      functionType: FunctionType.PUBLIC,
+      isInternal: true,
+      isStatic: false,
+      parameters: [
+        {
+          name: 'message_hash',
+          type: { kind: 'field' },
+          visibility: 'private' as ABIParameterVisibility,
+        },
+        {
+          name: 'authorize',
+          type: { kind: 'boolean' },
+          visibility: 'private' as ABIParameterVisibility,
+        },
+      ],
+      returnTypes: [],
+      errorTypes: {},
+    };
+  }
+}

package/src/utils/cross_chain.ts

@@ -0,0 +1,53 @@
+import type { Fr } from '@aztec/foundation/fields';
+import { retryUntil } from '@aztec/foundation/retry';
+import type { AztecNode } from '@aztec/stdlib/interfaces/client';
+
+/**
+ * Waits for the L1 to L2 message to be ready to be consumed.
+ * @param node - Aztec node instance used to obtain the information about the message
+ * @param l1ToL2MessageHash - Hash of the L1 to L2 message
+ * @param opts - Options
+ */
+export async function waitForL1ToL2MessageReady(
+  node: Pick<AztecNode, 'getBlockNumber' | 'getL1ToL2MessageBlock'>,
+  l1ToL2MessageHash: Fr,
+  opts: {
+    /** Timeout for the operation in seconds */ timeoutSeconds: number;
+    /** True if the message is meant to be consumed from a public function */ forPublicConsumption: boolean;
+  },
+) {
+  const messageBlockNumber = await node.getL1ToL2MessageBlock(l1ToL2MessageHash);
+  return retryUntil(
+    () => isL1ToL2MessageReady(node, l1ToL2MessageHash, { ...opts, messageBlockNumber }),
+    `L1 to L2 message ${l1ToL2MessageHash.toString()} ready`,
+    opts.timeoutSeconds,
+    1,
+  );
+}
+
+/**
+ * Returns whether the L1 to L2 message is ready to be consumed.
+ * @param node - Aztec node instance used to obtain the information about the message
+ * @param l1ToL2MessageHash - Hash of the L1 to L2 message
+ * @param opts - Options
+ * @returns True if the message is ready to be consumed, false otherwise
+ */
+export async function isL1ToL2MessageReady(
+  node: Pick<AztecNode, 'getBlockNumber' | 'getL1ToL2MessageBlock'>,
+  l1ToL2MessageHash: Fr,
+  opts: {
+    /** True if the message is meant to be consumed from a public function */ forPublicConsumption: boolean;
+    /** Cached synced block number for the message (will be fetched from PXE otherwise) */ messageBlockNumber?: number;
+  },
+): Promise<boolean> {
+  const blockNumber = await node.getBlockNumber();
+  const messageBlockNumber = opts.messageBlockNumber ?? (await node.getL1ToL2MessageBlock(l1ToL2MessageHash));
+  if (messageBlockNumber === undefined) {
+    return false;
+  }
+
+  // Note that public messages can be consumed 1 block earlier, since the sequencer will include the messages
+  // in the L1 to L2 message tree before executing the txs for the block. In private, however, we need to wait
+  // until the message is included so we can make use of the membership witness.
+  return opts.forPublicConsumption ? blockNumber + 1 >= messageBlockNumber : blockNumber >= messageBlockNumber;
+}

package/src/utils/fee_juice.ts

@@ -2,12 +2,13 @@ import { Fr } from '@aztec/foundation/fields';
 import { ProtocolContractAddress } from '@aztec/protocol-contracts';
 import type { AztecAddress } from '@aztec/stdlib/aztec-address';
 import { deriveStorageSlotInMap } from '@aztec/stdlib/hash';
-import type { PXE } from '@aztec/stdlib/interfaces/client';
+import type { AztecNode } from '@aztec/stdlib/interfaces/client';
 
 /**
  * Returns the owner's fee juice balance.
+ * Note: This is used only e2e_sandbox_example test. TODO: Consider nuking.
  */
-export async function getFeeJuiceBalance(owner: AztecAddress, pxe: PXE): Promise<bigint> {
+export async function getFeeJuiceBalance(owner: AztecAddress, node: AztecNode): Promise<bigint> {
   const slot = await deriveStorageSlotInMap(new Fr(1), owner);
-  return (await pxe.getPublicStorageAt(ProtocolContractAddress.FeeJuice, slot)).toBigInt();
+  return (await node.getPublicStorageAt('latest', ProtocolContractAddress.FeeJuice, slot)).toBigInt();
 }