added token 2022 support for token balance

packages/plugin-token/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@wardenprotocol/plugin-token",
-  "version": "2.0.15",
+  "version": "2.0.18",
   "main": "dist/index.cjs",
   "module": "dist/index.js",
   "types": "dist/index.d.ts",

packages/plugin-token/src/solana/tools/get_token_balances.ts

@@ -1,5 +1,5 @@
-import { TOKEN_PROGRAM_ID } from "@solana/spl-token";
-import { LAMPORTS_PER_SOL, type PublicKey } from "@solana/web3.js";
+import { TOKEN_2022_PROGRAM_ID, TOKEN_PROGRAM_ID } from "@solana/spl-token";
+import { LAMPORTS_PER_SOL, PublicKey } from "@solana/web3.js";
 import { SolanaAgentKit } from "solana-agent-kit";
 import { getTokenMetadata } from "./utils/tokenMetadata";
 
@@ -22,44 +22,60 @@ export async function get_token_balance(
     decimals: number;
   }>;
 }> {
-  const [lamportsBalance, tokenAccountData] = await Promise.all([
-    agent.connection.getBalance(walletAddress ?? agent.wallet.publicKey),
-    agent.connection.getParsedTokenAccountsByOwner(
-      walletAddress ?? agent.wallet.publicKey,
-      {
+  const owner = walletAddress ?? agent.wallet.publicKey;
+  const [lamportsBalance, tokenAccountData, token2022AccountData] =
+    await Promise.all([
+      agent.connection.getBalance(owner),
+      agent.connection.getParsedTokenAccountsByOwner(owner, {
         programId: TOKEN_PROGRAM_ID,
-      },
-    ),
-  ]);
+      }),
+      agent.connection.getParsedTokenAccountsByOwner(owner, {
+        programId: TOKEN_2022_PROGRAM_ID,
+      }),
+    ]);
 
-  const removedZeroBalance = tokenAccountData.value.filter(
-    (v: any) => v.account.data.parsed.info.tokenAmount.uiAmount !== 0,
-  );
+  // Process tokens from both programs separately since we know which program each belongs to
+  const processTokens = async (tokens: any[], isToken2022: boolean) => {
+    const nonZeroTokens = tokens.filter(
+      (v: any) => v.account.data.parsed.info.tokenAmount.uiAmount !== 0,
+    );
 
-  const tokenBalances = await Promise.all(
-    removedZeroBalance.map(async (v: any) => {
-      const mint = v.account.data.parsed.info.mint;
-      const mintInfo = await getTokenMetadata(agent.connection, mint);
-      // If metadata is not found, return Null values to not let the tool fail
-      if (!mintInfo) {
+    return Promise.all(
+      nonZeroTokens.map(async (v: any) => {
+        const mint = v.account.data.parsed.info.mint;
+        const mintInfo = await getTokenMetadata(
+          agent.connection,
+          mint,
+          isToken2022,
+        );
+        // If metadata is not found, return empty string values to not let the tool fail
+        if (!mintInfo) {
+          return {
+            tokenAddress: mint,
+            name: "",
+            symbol: "",
+            balance: v.account.data.parsed.info.tokenAmount.uiAmount as number,
+            decimals: v.account.data.parsed.info.tokenAmount.decimals as number,
+          };
+        }
+        // Normal return
         return {
           tokenAddress: mint,
-          name: "NULL",
-          symbol: "NULL",
+          name: mintInfo.name ?? "",
+          symbol: mintInfo.symbol ?? "",
           balance: v.account.data.parsed.info.tokenAmount.uiAmount as number,
           decimals: v.account.data.parsed.info.tokenAmount.decimals as number,
         };
-      }
-      // Normal return
-      return {
-        tokenAddress: mint,
-        name: mintInfo.name ?? "",
-        symbol: mintInfo.symbol ?? "",
-        balance: v.account.data.parsed.info.tokenAmount.uiAmount as number,
-        decimals: v.account.data.parsed.info.tokenAmount.decimals as number,
-      };
-    }),
-  );
+      }),
+    );
+  };
+
+  const [legacyTokenBalances, token2022Balances] = await Promise.all([
+    processTokens(tokenAccountData.value, false),
+    processTokens(token2022AccountData.value, true),
+  ]);
+
+  const tokenBalances = [...legacyTokenBalances, ...token2022Balances];
 
   const solBalance = lamportsBalance / LAMPORTS_PER_SOL;
 

packages/plugin-token/src/solana/tools/transfer.ts

@@ -1,6 +1,6 @@
 import {
   createAssociatedTokenAccountInstruction,
-  createTransferInstruction,
+  createTransferCheckedInstruction,
   getAccount,
   getAssociatedTokenAddress,
   getMint,
@@ -42,15 +42,29 @@ export async function transfer(
       tx = await signOrSendTX(agent, transaction.instructions);
     } else {
       const transaction = new Transaction();
+
+      const mintAccountInfo = await agent.connection.getAccountInfo(mint);
+      if (!mintAccountInfo) {
+        throw new Error("Mint account not found");
+      }
+      const tokenProgramId = mintAccountInfo.owner;
+
       // Transfer SPL token
       const fromAta = await getAssociatedTokenAddress(
         mint,
         agent.wallet.publicKey,
+        false,
+        tokenProgramId,
+      );
+      const toAta = await getAssociatedTokenAddress(
+        mint,
+        to,
+        false,
+        tokenProgramId,
       );
-      const toAta = await getAssociatedTokenAddress(mint, to);
 
       try {
-        await getAccount(agent.connection, toAta);
+        await getAccount(agent.connection, toAta, undefined, tokenProgramId);
       } catch {
         // Error is thrown if the tokenAccount doesn't exist
         transaction.add(
@@ -59,20 +73,30 @@ export async function transfer(
             toAta,
             to,
             mint,
+            tokenProgramId,
           ),
         );
       }
 
       // Get mint info to determine decimals
-      const mintInfo = await getMint(agent.connection, mint);
-      const adjustedAmount = amount * Math.pow(10, mintInfo.decimals);
+      const mintInfo = await getMint(
+        agent.connection,
+        mint,
+        undefined,
+        tokenProgramId,
+      );
+      const adjustedAmount = BigInt(Math.round(amount * Math.pow(10, mintInfo.decimals)));
 
       transaction.add(
-        createTransferInstruction(
+        createTransferCheckedInstruction(
           fromAta,
+          mint,
           toAta,
           agent.wallet.publicKey,
           adjustedAmount,
+          mintInfo.decimals,
+          [],
+          tokenProgramId,
         ),
       );