fix: wait for accounts to catch up with notes when deployed

yarn-project/aztec.js/src/account/manager/deploy_account_sent_tx.ts

@@ -1,7 +1,9 @@
 import { FieldsOf } from '@aztec/circuits.js';
 import { TxHash, TxReceipt } from '@aztec/types';
 
-import { SentTx, WaitOpts, Wallet } from '../../index.js';
+import { DefaultWaitOpts, SentTx, WaitOpts } from '../../contract/sent_tx.js';
+import { Wallet } from '../../wallet/index.js';
+import { waitForAccountSynch } from './util.js';
 
 /** Extends a transaction receipt with a wallet instance for the newly deployed contract. */
 export type DeployAccountTxReceipt = FieldsOf<TxReceipt> & {
@@ -32,8 +34,9 @@ export class DeployAccountSentTx extends SentTx {
    * @param opts - Options for configuring the waiting for the tx to be mined.
    * @returns The transaction receipt with the wallet for the deployed account contract.
    */
-  public async wait(opts?: WaitOpts): Promise<DeployAccountTxReceipt> {
+  public async wait(opts: WaitOpts = DefaultWaitOpts): Promise<DeployAccountTxReceipt> {
     const receipt = await super.wait(opts);
+    await waitForAccountSynch(this.pxe, this.wallet.getCompleteAddress(), opts);
     return { ...receipt, wallet: this.wallet };
   }
 }

yarn-project/aztec.js/src/account/manager/index.ts

@@ -2,6 +2,7 @@ import { PublicKey, getContractDeploymentInfo } from '@aztec/circuits.js';
 import { Fr } from '@aztec/foundation/fields';
 import { CompleteAddress, GrumpkinPrivateKey, PXE } from '@aztec/types';
 
+import { DefaultWaitOpts } from '../../contract/sent_tx.js';
 import {
   AccountWalletWithPrivateKey,
   ContractDeployer,
@@ -12,6 +13,7 @@ import {
 import { AccountContract, Salt } from '../index.js';
 import { AccountInterface } from '../interface.js';
 import { DeployAccountSentTx } from './deploy_account_sent_tx.js';
+import { waitForAccountSynch } from './util.js';
 
 /**
  * Manages a user account. Provides methods for calculating the account's address, deploying the account contract,
@@ -88,11 +90,12 @@ export class AccountManager {
    * Registers this account in the PXE Service and returns the associated wallet. Registering
    * the account on the PXE Service is required for managing private state associated with it.
    * Use the returned wallet to create Contract instances to be interacted with from this account.
+   * @param opts - Options to wait for the account to be synched.
    * @returns A Wallet instance.
    */
-  public async register(): Promise<AccountWalletWithPrivateKey> {
-    const completeAddress = await this.getCompleteAddress();
-    await this.pxe.registerAccount(this.encryptionPrivateKey, completeAddress.partialAddress);
+  public async register(opts: WaitOpts = DefaultWaitOpts): Promise<AccountWalletWithPrivateKey> {
+    const address = await this.#register();
+    await waitForAccountSynch(this.pxe, address, opts);
     return this.getWallet();
   }
 
@@ -105,7 +108,7 @@ export class AccountManager {
   public async getDeployMethod() {
     if (!this.deployMethod) {
       if (!this.salt) throw new Error(`Cannot deploy account contract without known salt.`);
-      await this.register();
+      await this.#register();
       const encryptionPublicKey = await this.getEncryptionPublicKey();
       const deployer = new ContractDeployer(this.accountContract.getContractArtifact(), this.pxe, encryptionPublicKey);
       const args = await this.accountContract.getDeploymentArgs();
@@ -138,8 +141,14 @@ export class AccountManager {
    * @param opts - Options to wait for the tx to be mined.
    * @returns A Wallet instance.
    */
-  public async waitDeploy(opts: WaitOpts = {}): Promise<AccountWalletWithPrivateKey> {
+  public async waitDeploy(opts: WaitOpts = DefaultWaitOpts): Promise<AccountWalletWithPrivateKey> {
     await this.deploy().then(tx => tx.wait(opts));
     return this.getWallet();
   }
+
+  async #register(): Promise<CompleteAddress> {
+    const completeAddress = await this.getCompleteAddress();
+    await this.pxe.registerAccount(this.encryptionPrivateKey, completeAddress.partialAddress);
+    return completeAddress;
+  }
 }

yarn-project/aztec.js/src/account/manager/util.ts

@@ -0,0 +1,29 @@
+import { CompleteAddress, PXE, WaitOpts, retryUntil } from '../../index.js';
+
+/**
+ * Waits for the account to finish synchronizing with the PXE Service.
+ * @param pxe - PXE instance
+ * @param address - Address to wait for synch
+ * @param opts - Wait options
+ */
+export async function waitForAccountSynch(
+  pxe: PXE,
+  address: CompleteAddress,
+  { interval, timeout }: WaitOpts,
+): Promise<void> {
+  const publicKey = address.publicKey.toString();
+  await retryUntil(
+    async () => {
+      const status = await pxe.getSyncStatus();
+      const accountSynchedToBlock = status.notes[publicKey];
+      if (typeof accountSynchedToBlock === 'undefined') {
+        return false;
+      } else {
+        return accountSynchedToBlock >= status.blocks;
+      }
+    },
+    'waitForAccountSynch',
+    timeout,
+    interval,
+  );
+}

yarn-project/aztec.js/src/contract/index.ts

@@ -37,6 +37,6 @@
  */
 export * from './contract.js';
 export * from './contract_function_interaction.js';
-export * from './sent_tx.js';
+export { SentTx, WaitOpts } from './sent_tx.js';
 export * from './contract_base.js';
 export * from './batch_call.js';

yarn-project/aztec.js/src/contract/sent_tx.ts

@@ -19,7 +19,7 @@ export type WaitOpts = {
   debug?: boolean;
 };
 
-const DefaultWaitOpts: WaitOpts = {
+export const DefaultWaitOpts: WaitOpts = {
   timeout: 60,
   interval: 1,
   waitForNotesSync: true,

yarn-project/end-to-end/src/e2e_2_pxes.test.ts

@@ -6,11 +6,13 @@ import {
   EthAddress,
   ExtendedNote,
   Fr,
+  GrumpkinScalar,
   Note,
   PXE,
   TxStatus,
   Wallet,
   computeMessageSecretHash,
+  getUnsafeSchnorrAccount,
   retryUntil,
 } from '@aztec/aztec.js';
 import { ChildContract, TokenContract } from '@aztec/noir-contracts/types';
@@ -248,4 +250,22 @@ describe('e2e_2_pxes', () => {
     // Check that user B balance is 0 on server A
     await expectTokenBalance(walletA, completeTokenAddress.address, userB.address, 0n, checkIfSynchronized);
   });
+
+  it('permits migrating an account from one PXE to another', async () => {
+    const privateKey = GrumpkinScalar.random();
+    const account = getUnsafeSchnorrAccount(pxeA, privateKey, Fr.random());
+    const completeAddress = await account.getCompleteAddress();
+    const wallet = await account.waitDeploy();
+
+    await expect(wallet.isAccountStateSynchronized(completeAddress.address)).resolves.toBe(true);
+    const accountOnB = getUnsafeSchnorrAccount(pxeB, privateKey, completeAddress);
+    const walletOnB = await accountOnB.getWallet();
+
+    // need to register first otherwise the new PXE won't know about the account
+    await expect(walletOnB.isAccountStateSynchronized(completeAddress.address)).rejects.toThrow();
+
+    await accountOnB.register();
+    // registering should wait for the account to be synchronized
+    await expect(walletOnB.isAccountStateSynchronized(completeAddress.address)).resolves.toBe(true);
+  });
 });

yarn-project/pxe/src/synchronizer/synchronizer.ts

@@ -240,7 +240,8 @@ export class Synchronizer {
    * @returns A promise that resolves once the account is added to the Synchronizer.
    */
   public addAccount(publicKey: PublicKey, keyStore: KeyStore, startingBlock: number) {
-    const processor = this.noteProcessors.find(x => x.publicKey.equals(publicKey));
+    const predicate = (x: NoteProcessor) => x.publicKey.equals(publicKey);
+    const processor = this.noteProcessors.find(predicate) ?? this.noteProcessorsToCatchUp.find(predicate);
     if (processor) return;
 
     this.noteProcessorsToCatchUp.push(new NoteProcessor(publicKey, keyStore, this.db, this.node, startingBlock));
@@ -259,7 +260,8 @@ export class Synchronizer {
     if (!completeAddress) {
       throw new Error(`Checking if account is synched is not possible for ${account} because it is not registered.`);
     }
-    const processor = this.noteProcessors.find(x => x.publicKey.equals(completeAddress.publicKey));
+    const findByPublicKey = (x: NoteProcessor) => x.publicKey.equals(completeAddress.publicKey);
+    const processor = this.noteProcessors.find(findByPublicKey) ?? this.noteProcessorsToCatchUp.find(findByPublicKey);
     if (!processor) {
       throw new Error(
         `Checking if account is synched is not possible for ${account} because it is only registered as a recipient.`,