From 977658bcb7bb376fe528c69e2240a77cdb0bd1c3 Mon Sep 17 00:00:00 2001
From: Yolley <Yolley@users.noreply.github.com>
Date: Mon, 25 Dec 2023 21:24:31 +0900
Subject: [PATCH] create ata on withdrawal if it does not exist (#115)

---
 lerna.json                             |  2 +-
 packages/stream/package.json           |  2 +-
 packages/stream/solana/StreamClient.ts | 27 ++++++++++++++++++++++++++
 packages/stream/solana/types.ts        |  6 ++++++
 4 files changed, 35 insertions(+), 2 deletions(-)

diff --git a/lerna.json b/lerna.json
index 705c2f62..18865450 100644
--- a/lerna.json
+++ b/lerna.json
@@ -2,6 +2,6 @@
   "packages": [
     "packages/*"
   ],
-  "version": "5.8.1",
+  "version": "5.9.0",
   "$schema": "node_modules/lerna/schemas/lerna-schema.json"
 }
\ No newline at end of file
diff --git a/packages/stream/package.json b/packages/stream/package.json
index 3a3d5684..84a10674 100644
--- a/packages/stream/package.json
+++ b/packages/stream/package.json
@@ -1,6 +1,6 @@
 {
   "name": "@streamflow/stream",
-  "version": "5.8.1",
+  "version": "5.9.0",
   "description": "JavaScript SDK to interact with Streamflow protocol.",
   "main": "dist/index.js",
   "homepage": "https://github.com/streamflow-finance/js-sdk/",
diff --git a/packages/stream/solana/StreamClient.ts b/packages/stream/solana/StreamClient.ts
index bcf0988e..c3c5f3f8 100644
--- a/packages/stream/solana/StreamClient.ts
+++ b/packages/stream/solana/StreamClient.ts
@@ -18,6 +18,7 @@ import {
   sendAndConfirmRawTransaction,
   BlockheightBasedTransactionConfirmationStrategy,
 } from "@solana/web3.js";
+import { createAssociatedTokenAccountInstruction } from "@solana/spl-token";
 import * as borsh from "borsh";
 
 import {
@@ -30,6 +31,7 @@ import {
   ICreateStreamSolanaExt,
   IInteractStreamSolanaExt,
   ITopUpStreamSolanaExt,
+  CheckAssociatedTokenAccountData,
 } from "./types";
 import {
   ata,
@@ -477,6 +479,7 @@ export default class SolanaStreamClient extends BaseStreamClient {
 
     const streamflowTreasuryTokens = await ata(data.mint, STREAMFLOW_TREASURY_PUBLIC_KEY);
     const partnerTokens = await ata(data.mint, data.partner);
+    await this.checkAssociatedTokenAccount(data, { invoker }, ixs);
 
     ixs.push(
       withdrawStreamInstruction(amount, this.programId, {
@@ -531,6 +534,7 @@ export default class SolanaStreamClient extends BaseStreamClient {
     const partnerTokens = await ata(data.mint, data.partner);
 
     const ixs: TransactionInstruction[] = [];
+    await this.checkAssociatedTokenAccount(data, { invoker }, ixs);
 
     ixs.push(
       cancelStreamInstruction(this.programId, {
@@ -942,6 +946,29 @@ export default class SolanaStreamClient extends BaseStreamClient {
     return { tx, metadataPubKey };
   }
 
+  /**
+   * Utility function that checks whether associated token account for the recipient exists and adds an instruction to add if not
+   */
+  private async checkAssociatedTokenAccount(
+    data: CheckAssociatedTokenAccountData,
+    { invoker }: IInteractStreamSolanaExt,
+    ixs: TransactionInstruction[]
+  ) {
+    const accountExists = await this.connection.getAccountInfo(data.recipientTokens);
+    if (!accountExists?.data) {
+      ixs.push(
+        createAssociatedTokenAccountInstruction(
+          invoker.publicKey!,
+          data.recipientTokens,
+          data.recipient,
+          data.mint,
+          TOKEN_PROGRAM_ID,
+          ASSOCIATED_TOKEN_PROGRAM_ID
+        )
+      );
+    }
+  }
+
   /**
    * Utility function to generate metadata for a Contract or return existing Pubkey
    */
diff --git a/packages/stream/solana/types.ts b/packages/stream/solana/types.ts
index aec06878..277f786a 100644
--- a/packages/stream/solana/types.ts
+++ b/packages/stream/solana/types.ts
@@ -426,6 +426,12 @@ export interface CreateMultipleStreamsValues {
   recipients: Recipient[];
 }
 
+export interface CheckAssociatedTokenAccountData {
+  recipient: PublicKey;
+  recipientTokens: PublicKey;
+  mint: PublicKey;
+}
+
 export interface BatchItem {
   recipient: string;
   tx: Transaction;