refactor(sdk): fix async handling

prelim tests passing
Squads-Protocol · Oct 19, 2024 · 50ca3b7 · 50ca3b7
1 parent ec367d2
commit 50ca3b7
Show file tree

Hide file tree

Showing 9 changed files with 241 additions and 174 deletions.
diff --git a/sdk/multisig/src/actions/common.ts b/sdk/multisig/src/actions/common.ts
@@ -1,5 +1,6 @@
 import {
   Connection,
+  Keypair,
   PublicKey,
   SendOptions,
   Signer,
@@ -31,25 +32,39 @@ export abstract class BaseBuilder<
   T extends BuildResult,
   U extends BaseBuilderArgs = BaseBuilderArgs
 > {
+  public createKey: Keypair;
   protected connection: Connection;
   protected instructions: TransactionInstruction[] = [];
   protected creator: PublicKey = PublicKey.default;
-  protected buildPromise: Promise<T>;
+  protected buildPromise: Promise<void>;
   protected args: Omit<U, keyof BaseBuilderArgs>;
+  private built: boolean = false;
 
   constructor(args: U) {
     this.connection = args.connection;
     this.creator = args.creator;
     this.args = this.extractAdditionalArgs(args);
-    this.buildPromise = this.build();
+    this.createKey = Keypair.generate();
+    this.buildPromise = this.initializeBuild();
+  }
+
+  private async initializeBuild(): Promise<void> {
+    await this.build();
+    this.built = true;
+  }
+
+  protected async ensureBuilt(): Promise<void> {
+    if (!this.built) {
+      await this.buildPromise;
+    }
   }
 
   private extractAdditionalArgs(args: U): Omit<U, keyof BaseBuilderArgs> {
     const { connection, creator, ...additionalArgs } = args;
     return additionalArgs;
   }
 
-  protected abstract build(): Promise<T>;
+  protected abstract build(): Promise<void>;
 
   getInstructions(): TransactionInstruction[] {
     return this.instructions;
@@ -76,19 +91,18 @@ export abstract class BaseBuilder<
    *
    */
   async transaction(feePayer?: Signer): Promise<VersionedTransaction> {
-    return this.buildPromise.then(async () => {
-      const message = new TransactionMessage({
-        payerKey: feePayer?.publicKey ?? this.creator!,
-        recentBlockhash: (await this.connection.getLatestBlockhash()).blockhash,
-        instructions: [...this.instructions!],
-      }).compileToV0Message();
-
-      const tx = new VersionedTransaction(message);
-      if (feePayer) {
-        tx.sign([feePayer]);
-      }
-      return tx;
-    });
+    await this.ensureBuilt();
+    const message = new TransactionMessage({
+      payerKey: feePayer?.publicKey ?? this.creator,
+      recentBlockhash: (await this.connection.getLatestBlockhash()).blockhash,
+      instructions: [...this.instructions],
+    }).compileToV0Message();
+
+    const tx = new VersionedTransaction(message);
+    if (feePayer) {
+      tx.sign([feePayer]);
+    }
+    return tx;
   }
 
   /**
@@ -103,32 +117,35 @@ export abstract class BaseBuilder<
     preInstructions?: TransactionInstruction[];
     postInstructions?: TransactionInstruction[];
     feePayer?: Signer;
+    signers?: Signer[];
     options?: SendOptions;
   }): Promise<TransactionSignature> {
-    return this.buildPromise.then(async () => {
-      const instructions = [...this.instructions];
-      if (settings?.preInstructions) {
-        instructions.unshift(...settings.preInstructions);
-      }
-      if (settings?.postInstructions) {
-        instructions.push(...settings.postInstructions);
-      }
-      const message = new TransactionMessage({
-        payerKey: settings?.feePayer?.publicKey ?? this.creator,
-        recentBlockhash: (await this.connection.getLatestBlockhash()).blockhash,
-        instructions: [...instructions],
-      }).compileToV0Message();
-
-      const tx = new VersionedTransaction(message);
-      if (settings?.feePayer) {
-        tx.sign([settings.feePayer]);
-      }
-      const signature = await this.connection.sendTransaction(
-        tx,
-        settings?.options
-      );
-      return signature;
-    });
+    await this.ensureBuilt();
+    const instructions = [...this.instructions];
+    if (settings?.preInstructions) {
+      instructions.unshift(...settings.preInstructions);
+    }
+    if (settings?.postInstructions) {
+      instructions.push(...settings.postInstructions);
+    }
+    const message = new TransactionMessage({
+      payerKey: settings?.feePayer?.publicKey ?? this.creator,
+      recentBlockhash: (await this.connection.getLatestBlockhash()).blockhash,
+      instructions: [...instructions],
+    }).compileToV0Message();
+
+    const tx = new VersionedTransaction(message);
+    if (settings?.feePayer) {
+      tx.sign([settings.feePayer]);
+    }
+    if (settings?.signers) {
+      tx.sign([...settings.signers]);
+    }
+    const signature = await this.connection.sendTransaction(
+      tx,
+      settings?.options
+    );
+    return signature;
   }
 
   /**
@@ -143,33 +160,50 @@ export abstract class BaseBuilder<
     preInstructions?: TransactionInstruction[];
     postInstructions?: TransactionInstruction[];
     feePayer?: Signer;
+    signers?: Signer[];
     options?: SendOptions;
   }): Promise<TransactionSignature> {
-    return this.buildPromise.then(async () => {
-      const instructions = [...this.instructions];
-      if (settings?.preInstructions) {
-        instructions.unshift(...settings.preInstructions);
-      }
-      if (settings?.postInstructions) {
-        instructions.push(...settings.postInstructions);
-      }
-      const message = new TransactionMessage({
-        payerKey: settings?.feePayer?.publicKey ?? this.creator,
-        recentBlockhash: (await this.connection.getLatestBlockhash()).blockhash,
-        instructions: [...instructions],
-      }).compileToV0Message();
-
-      const tx = new VersionedTransaction(message);
-      if (settings?.feePayer) {
-        tx.sign([settings.feePayer]);
+    await this.ensureBuilt();
+    const instructions = [...this.instructions];
+    if (settings?.preInstructions) {
+      instructions.unshift(...settings.preInstructions);
+    }
+    if (settings?.postInstructions) {
+      instructions.push(...settings.postInstructions);
+    }
+    const message = new TransactionMessage({
+      payerKey: settings?.feePayer?.publicKey ?? this.creator,
+      recentBlockhash: (await this.connection.getLatestBlockhash()).blockhash,
+      instructions: [...instructions],
+    }).compileToV0Message();
+
+    const tx = new VersionedTransaction(message);
+    if (settings?.feePayer) {
+      tx.sign([settings.feePayer]);
+    }
+    if (settings?.signers) {
+      tx.sign([...settings.signers]);
+    }
+    const signature = await this.connection.sendTransaction(
+      tx,
+      settings?.options
+    );
+
+    let commitment = settings?.options?.preflightCommitment;
+
+    let sent = false;
+    while (sent === false) {
+      const status = await this.connection.getSignatureStatuses([signature]);
+      if (
+        commitment
+          ? status.value[0]?.confirmationStatus === commitment
+          : status.value[0]?.confirmationStatus === "finalized"
+      ) {
+        sent = true;
       }
-      const signature = await this.connection.sendTransaction(
-        tx,
-        settings?.options
-      );
-      await this.connection.getSignatureStatuses([signature]);
-      return signature;
-    });
+    }
+
+    return signature;
   }
 
   /**
@@ -201,24 +235,25 @@ export abstract class BaseBuilder<
   async customSend(
     callback: (args: TransactionMessage) => Promise<TransactionSignature>
   ): Promise<TransactionSignature> {
-    return this.buildPromise.then(async () => {
-      const message = new TransactionMessage({
-        payerKey: this.creator,
-        recentBlockhash: (await this.connection.getLatestBlockhash()).blockhash,
-        instructions: [...this.instructions],
-      });
-
-      const signature = await callback(message);
-      return signature;
+    await this.ensureBuilt();
+    const message = new TransactionMessage({
+      payerKey: this.creator,
+      recentBlockhash: (await this.connection.getLatestBlockhash()).blockhash,
+      instructions: [...this.instructions],
     });
+
+    const signature = await callback(message);
+    return signature;
   }
 
+  /*
   protected then<TResult1 = T, TResult2 = never>(
     onfulfilled?: ((value: T) => TResult1 | PromiseLike<TResult1>) | null,
     onrejected?: ((reason: any) => TResult2 | PromiseLike<TResult2>) | null
   ): Promise<TResult1 | TResult2> {
     return this.buildPromise.then(onfulfilled, onrejected);
   }
+  */
 }
 
 export interface TransactionBuilderArgs extends BaseBuilderArgs {

diff --git a/sdk/multisig/src/actions/createBatch.ts b/sdk/multisig/src/actions/createBatch.ts
@@ -135,13 +135,13 @@ class BatchBuilder extends BaseBuilder<
 > {
   public instructions: TransactionInstruction[] = [];
   public index: number = 1;
-  public innerIndex: number = 0;
+  public innerIndex: number = 1;
 
   constructor(args: CreateBatchActionArgs) {
     super(args);
   }
 
-  async build() {
+  protected async build() {
     const {
       multisig,
       vaultIndex = 0,
@@ -161,18 +161,10 @@ class BatchBuilder extends BaseBuilder<
 
     this.instructions = [...result.instructions];
     this.index = result.index;
-    return this;
-  }
-
-  getInstructions(): TransactionInstruction[] {
-    return this.instructions;
-  }
-
-  getIndex(): number {
-    return this.index;
   }
 
-  getBatchKey(): PublicKey | null {
+  async getBatchKey(): Promise<PublicKey> {
+    this.ensureBuilt();
     const index = this.index;
     const [batchPda] = getTransactionPda({
       multisigPda: this.args.multisig,
@@ -182,21 +174,23 @@ class BatchBuilder extends BaseBuilder<
     return batchPda;
   }
 
-  getBatchTransactionKey(innerIndex: number): PublicKey | null {
+  async getBatchTransactionKey(innerIndex?: number): Promise<PublicKey> {
+    this.ensureBuilt();
     const index = this.index;
     const [batchPda] = getBatchTransactionPda({
       multisigPda: this.args.multisig,
       batchIndex: BigInt(index ?? 1),
-      transactionIndex: this.innerIndex ?? 1,
+      transactionIndex: innerIndex ?? this.innerIndex,
     });
 
     return batchPda;
   }
 
-  getAllBatchTransactionKeys(localIndex: number): PublicKey[] | null {
+  async getAllBatchTransactionKeys(): Promise<PublicKey[]> {
+    this.ensureBuilt();
     const index = this.index;
     const transactions = [];
-    for (let i = 1; i <= localIndex; i++) {
+    for (let i = 1; i <= this.innerIndex; i++) {
       const [batchPda] = getBatchTransactionPda({
         multisigPda: this.args.multisig,
         batchIndex: BigInt(index ?? 1),
@@ -212,11 +206,10 @@ class BatchBuilder extends BaseBuilder<
   async getBatchAccount(
     key: PublicKey
   ): Promise<Pick<BatchBuilder, BatchMethods<"getBatchAccount">>> {
-    return this.buildPromise.then(async () => {
-      const batchAccount = await Batch.fromAccountAddress(this.connection, key);
+    this.ensureBuilt();
+    const batchAccount = await Batch.fromAccountAddress(this.connection, key);
 
-      return batchAccount;
-    });
+    return batchAccount;
   }
 
   /**
@@ -228,7 +221,7 @@ class BatchBuilder extends BaseBuilder<
     message: TransactionMessage,
     member?: PublicKey
   ): Promise<Pick<BatchBuilder, BatchMethods<"addTransaction">>> {
-    this.innerIndex++;
+    this.ensureBuilt();
     const { instruction } = await addBatchTransactionCore({
       multisig: this.args.multisig,
       member: member ?? this.creator,
@@ -242,6 +235,8 @@ class BatchBuilder extends BaseBuilder<
 
     this.instructions.push(instruction);
 
+    this.innerIndex++;
+
     return this;
   }
 
@@ -250,9 +245,10 @@ class BatchBuilder extends BaseBuilder<
    * @args feePayer - Optional signer to pay the transaction fee.
    * @returns `VersionedTransaction` with the `vaultTransactionCreate` instruction.
    */
-  withProposal(
+  async withProposal(
     isDraft?: boolean
-  ): Pick<BatchBuilder, BatchMethods<"withProposal">> {
+  ): Promise<Pick<BatchBuilder, BatchMethods<"withProposal">>> {
+    this.ensureBuilt();
     const { instruction } = createProposalCore({
       multisig: this.args.multisig,
       creator: this.creator,
@@ -314,6 +310,7 @@ class BatchBuilder extends BaseBuilder<
   async withExecute(
     member?: PublicKey
   ): Promise<Pick<BatchBuilder, BatchMethods<"withExecute">>> {
+    await this.ensureBuilt();
     const { instruction } = await executeBatchTransactionCore({
       connection: this.connection,
       multisig: this.args.multisig,

diff --git a/sdk/multisig/src/actions/createConfigTransaction.ts b/sdk/multisig/src/actions/createConfigTransaction.ts
@@ -117,17 +117,17 @@ class ConfigTransactionBuilder extends BaseTransactionBuilder<
 
     this.instructions = [...result.instructions];
     this.index = result.index;
-    return this;
   }
 
   /**
    * Creates a transaction containing the ConfigTransaction creation instruction.
    * @args feePayer - Optional signer to pay the transaction fee.
    * @returns `VersionedTransaction` with the `vaultTransactionCreate` instruction.
    */
-  withProposal(
+  async withProposal(
     isDraft?: boolean
-  ): Pick<ConfigTransactionBuilder, Methods<"withProposal">> {
+  ): Promise<Pick<ConfigTransactionBuilder, Methods<"withProposal">>> {
+    await this.ensureBuilt();
     const { instruction } = createProposalCore({
       multisig: this.args.multisig,
       creator: this.creator,