From f77e65579da8f8ebdbf0fe92c0f451e86da9a788 Mon Sep 17 00:00:00 2001
From: Santiago Palladino <santiago@aztecprotocol.com>
Date: Tue, 13 Aug 2024 15:11:48 -0300
Subject: [PATCH] feat: Add proven flag to sent tx wait opts

Allows to wait for a tx to be proven instead of it just being mined from
aztec.js sent-tx class.
---
 .../aztec.js/src/contract/sent_tx.test.ts     | 10 ++++++++++
 yarn-project/aztec.js/src/contract/sent_tx.ts | 20 +++++++++++++++++++
 2 files changed, 30 insertions(+)

diff --git a/yarn-project/aztec.js/src/contract/sent_tx.test.ts b/yarn-project/aztec.js/src/contract/sent_tx.test.ts
index 119fc766e60..d7125dd4755 100644
--- a/yarn-project/aztec.js/src/contract/sent_tx.test.ts
+++ b/yarn-project/aztec.js/src/contract/sent_tx.test.ts
@@ -48,5 +48,15 @@ describe('SentTx', () => {
       pxe.getSyncStatus.mockResolvedValue({ blocks: 19, notes: { '0x1': 19, '0x2': 19 } });
       await expect(sentTx.wait({ timeout: 1, interval: 0.4 })).rejects.toThrow(/dropped/);
     });
+
+    it('waits for the tx to be proven', async () => {
+      const waitOpts = { timeout: 1, interval: 0.4, waitForNotesSync: false, proven: true, provenTimeout: 2 };
+      pxe.getProvenBlockNumber.mockResolvedValue(10);
+      await expect(sentTx.wait(waitOpts)).rejects.toThrow(/timeout/i);
+
+      pxe.getProvenBlockNumber.mockResolvedValue(20);
+      const actual = await sentTx.wait(waitOpts);
+      expect(actual).toEqual(txReceipt);
+    });
   });
 });
diff --git a/yarn-project/aztec.js/src/contract/sent_tx.ts b/yarn-project/aztec.js/src/contract/sent_tx.ts
index 03916f83c44..91b0f5ea2c0 100644
--- a/yarn-project/aztec.js/src/contract/sent_tx.ts
+++ b/yarn-project/aztec.js/src/contract/sent_tx.ts
@@ -13,8 +13,12 @@ import { type FieldsOf } from '@aztec/foundation/types';
 export type WaitOpts = {
   /** The maximum time (in seconds) to wait for the transaction to be mined. Defaults to 60. */
   timeout?: number;
+  /** The maximum time (in seconds) to wait for the transaction to be proven. Defaults to 600. */
+  provenTimeout?: number;
   /** The time interval (in seconds) between retries to fetch the transaction receipt. Defaults to 1. */
   interval?: number;
+  /** Whether to wait for the tx to be proven. */
+  proven?: boolean;
   /**
    * Whether to wait for the PXE Service to sync all notes up to the block in which this tx was mined.
    * If false, then any queries that depend on state set by this transaction may return stale data. Defaults to true.
@@ -28,6 +32,7 @@ export type WaitOpts = {
 
 export const DefaultWaitOpts: WaitOpts = {
   timeout: 60,
+  provenTimeout: 600,
   interval: 1,
   waitForNotesSync: true,
   debug: false,
@@ -78,6 +83,9 @@ export class SentTx {
         `Transaction ${await this.getTxHash()} was ${receipt.status}. Reason: ${receipt.error ?? 'unknown'}`,
       );
     }
+    if (opts?.proven && receipt.blockNumber !== undefined) {
+      await this.waitForProven(receipt.blockNumber, opts);
+    }
     if (opts?.debug) {
       const txHash = await this.getTxHash();
       const tx = (await this.pxe.getTxEffect(txHash))!;
@@ -144,4 +152,16 @@ export class SentTx {
       opts?.interval ?? DefaultWaitOpts.interval,
     );
   }
+
+  protected async waitForProven(minedBlock: number, opts?: WaitOpts) {
+    return await retryUntil(
+      async () => {
+        const provenBlock = await this.pxe.getProvenBlockNumber();
+        return provenBlock >= minedBlock ? provenBlock : undefined;
+      },
+      'isProven',
+      opts?.provenTimeout ?? DefaultWaitOpts.provenTimeout,
+      opts?.interval ?? DefaultWaitOpts.interval,
+    );
+  }
 }