feat: get all token balances

src/actions/balance.ts

@@ -54,7 +54,6 @@ const balanceAction: Action = {
     return {
       status: "success",
       balance: balance,
-      token: input.tokenAddress || "SOL",
     };
   },
 };

src/actions/createOrcaSingleSidedWhirlpool.ts

@@ -87,7 +87,7 @@ const createOrcaSingleSidedWhirlpoolAction: Action = {
       const otherTokenMint = new PublicKey(input.otherTokenMint);
       const initialPrice = new Decimal(input.initialPrice);
       const maxPrice = new Decimal(input.maxPrice);
-      const feeTier = input.feeTier
+      const feeTier = input.feeTier;
 
       // Create the whirlpool
       const signature = await orcaCreateSingleSidedLiquidityPool(

src/actions/index.ts

@@ -28,34 +28,35 @@ import createOrcaSingleSidedWhirlpoolAction from "./createOrcaSingleSidedWhirlpo
 import launchPumpfunTokenAction from "./launchPumpfunToken";
 
 export const ACTIONS = {
-  "DEPLOY_TOKEN_ACTION" : deployTokenAction,
-  "BALANCE_ACTION" : balanceAction,
-  "TRANSFER_ACTION" : transferAction,
-  "DEPLOY_COLLECTION_ACTION" : deployCollectionAction,
-  "MINT_NFT_ACTION" : mintNFTAction,
-  "TRADE_ACTION" : tradeAction,
-  "REQUEST_FUNDS_ACTION" : requestFundsAction,
-  "RESOLVE_DOMAIN_ACTION" : resolveDomainAction,
-  "GET_TOKEN_DATA_ACTION" : getTokenDataAction,
-  "GET_TPS_ACTION" : getTPSAction,
-  "FETCH_PRICE_ACTION" : fetchPriceAction,
-  "STAKE_WITH_JUP_ACTION" : stakeWithJupAction,
-  "REGISTER_DOMAIN_ACTION" : registerDomainAction,
-  "LEND_ASSET_ACTION" : lendAssetAction,
-  "CREATE_GIBWORK_TASK_ACTION" : createGibworkTaskAction,
-  "RESOLVE_SOL_DOMAIN_ACTION" : resolveSolDomainAction,
-  "PYTH_FETCH_PRICE_ACTION" : pythFetchPriceAction,
-  "GET_OWNED_DOMAINS_FOR_TLD_ACTION" : getOwnedDomainsForTLDAction,
-  "GET_PRIMARY_DOMAIN_ACTION" : getPrimaryDomainAction,
-  "GET_ALL_DOMAINS_TLDS_ACTION" : getAllDomainsTLDsAction,
-  "GET_OWNED_ALL_DOMAINS_ACTION" : getOwnedAllDomainsAction,
-  "CREATE_IMAGE_ACTION" : createImageAction,
-  "GET_MAIN_ALL_DOMAINS_DOMAIN_ACTION" : getMainAllDomainsDomainAction,
-  "GET_ALL_REGISTERED_ALL_DOMAINS_ACTION" : getAllRegisteredAllDomainsAction,
-  "RAYDIUM_CREATE_CPMM_ACTION" : raydiumCreateCpmmAction,
-  "RAYDIUM_CREATE_AMM_V4_ACTION" : raydiumCreateAmmV4Action,
-  "CREATE_ORCA_SINGLE_SIDED_WHIRLPOOL_ACTION" : createOrcaSingleSidedWhirlpoolAction,
-  "LAUNCH_PUMPFUN_TOKEN_ACTION" : launchPumpfunTokenAction,
+  DEPLOY_TOKEN_ACTION: deployTokenAction,
+  BALANCE_ACTION: balanceAction,
+  TRANSFER_ACTION: transferAction,
+  DEPLOY_COLLECTION_ACTION: deployCollectionAction,
+  MINT_NFT_ACTION: mintNFTAction,
+  TRADE_ACTION: tradeAction,
+  REQUEST_FUNDS_ACTION: requestFundsAction,
+  RESOLVE_DOMAIN_ACTION: resolveDomainAction,
+  GET_TOKEN_DATA_ACTION: getTokenDataAction,
+  GET_TPS_ACTION: getTPSAction,
+  FETCH_PRICE_ACTION: fetchPriceAction,
+  STAKE_WITH_JUP_ACTION: stakeWithJupAction,
+  REGISTER_DOMAIN_ACTION: registerDomainAction,
+  LEND_ASSET_ACTION: lendAssetAction,
+  CREATE_GIBWORK_TASK_ACTION: createGibworkTaskAction,
+  RESOLVE_SOL_DOMAIN_ACTION: resolveSolDomainAction,
+  PYTH_FETCH_PRICE_ACTION: pythFetchPriceAction,
+  GET_OWNED_DOMAINS_FOR_TLD_ACTION: getOwnedDomainsForTLDAction,
+  GET_PRIMARY_DOMAIN_ACTION: getPrimaryDomainAction,
+  GET_ALL_DOMAINS_TLDS_ACTION: getAllDomainsTLDsAction,
+  GET_OWNED_ALL_DOMAINS_ACTION: getOwnedAllDomainsAction,
+  CREATE_IMAGE_ACTION: createImageAction,
+  GET_MAIN_ALL_DOMAINS_DOMAIN_ACTION: getMainAllDomainsDomainAction,
+  GET_ALL_REGISTERED_ALL_DOMAINS_ACTION: getAllRegisteredAllDomainsAction,
+  RAYDIUM_CREATE_CPMM_ACTION: raydiumCreateCpmmAction,
+  RAYDIUM_CREATE_AMM_V4_ACTION: raydiumCreateAmmV4Action,
+  CREATE_ORCA_SINGLE_SIDED_WHIRLPOOL_ACTION:
+    createOrcaSingleSidedWhirlpoolAction,
+  LAUNCH_PUMPFUN_TOKEN_ACTION: launchPumpfunTokenAction,
 };
 
 export type { Action, ActionExample, Handler } from "../types/action";

src/agent/index.ts

@@ -83,24 +83,30 @@ export class SolanaAgentKit {
    * @deprecated Using openai_api_key directly in constructor is deprecated.
    * Please use the new constructor with Config object instead:
    * @example
-   * const agent = new SolanaAgentKit(privateKey, rpcUrl, { 
+   * const agent = new SolanaAgentKit(privateKey, rpcUrl, {
    *   OPENAI_API_KEY: 'your-key'
    * });
    */
-  constructor(private_key: string, rpc_url: string, openai_api_key: string | null);
+  constructor(
+    private_key: string,
+    rpc_url: string,
+    openai_api_key: string | null,
+  );
   constructor(private_key: string, rpc_url: string, config: Config);
   constructor(
     private_key: string,
     rpc_url: string,
     configOrKey: Config | string | null,
   ) {
-    this.connection = new Connection(rpc_url || "https://api.mainnet-beta.solana.com");
+    this.connection = new Connection(
+      rpc_url || "https://api.mainnet-beta.solana.com",
+    );
     this.wallet = Keypair.fromSecretKey(bs58.decode(private_key));
     this.wallet_address = this.wallet.publicKey;
 
     // Handle both old and new patterns
-    if (typeof configOrKey === 'string' || configOrKey === null) {
-      this.config = { OPENAI_API_KEY: configOrKey || '' };
+    if (typeof configOrKey === "string" || configOrKey === null) {
+      this.config = { OPENAI_API_KEY: configOrKey || "" };
     } else {
       this.config = configOrKey;
     }
@@ -127,7 +133,19 @@ export class SolanaAgentKit {
     return deploy_collection(this, options);
   }
 
-  async getBalance(token_address?: PublicKey): Promise<number> {
+  async getBalance(token_address?: PublicKey): Promise<
+    | number
+    | {
+        sol: number;
+        tokens: Array<{
+          tokenAddress: string;
+          name: string;
+          symbol: string;
+          balance: number;
+          decimals: number;
+        }>;
+      }
+  > {
     return get_balance(this, token_address);
   }
 

src/tools/get_balance.ts

@@ -1,5 +1,7 @@
 import { LAMPORTS_PER_SOL, PublicKey } from "@solana/web3.js";
 import { SolanaAgentKit } from "../index";
+import { TOKEN_PROGRAM_ID } from "@solana/spl-token";
+import { getTokenMetadata } from "../utils/tokenMetadata";
 
 /**
  * Get the balance of SOL or an SPL token for the agent's wallet
@@ -10,12 +12,51 @@ import { SolanaAgentKit } from "../index";
 export async function get_balance(
   agent: SolanaAgentKit,
   token_address?: PublicKey,
-): Promise<number> {
+): Promise<
+  | number
+  | {
+      sol: number;
+      tokens: Array<{
+        tokenAddress: string;
+        name: string;
+        symbol: string;
+        balance: number;
+        decimals: number;
+      }>;
+    }
+> {
   if (!token_address) {
-    return (
-      (await agent.connection.getBalance(agent.wallet_address)) /
-      LAMPORTS_PER_SOL
+    const [lamportsBalance, tokenAccountData] = await Promise.all([
+      agent.connection.getBalance(agent.wallet_address),
+      agent.connection.getParsedTokenAccountsByOwner(agent.wallet_address, {
+        programId: TOKEN_PROGRAM_ID,
+      }),
+    ]);
+
+    const removedZeroBalance = tokenAccountData.value.filter(
+      (v) => v.account.data.parsed.info.tokenAmount.uiAmount !== 0,
     );
+
+    const tokenBalances = await Promise.all(
+      removedZeroBalance.map(async (v) => {
+        const mint = v.account.data.parsed.info.mint;
+        const mintInfo = await getTokenMetadata(agent.connection, mint);
+        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 solBalance = lamportsBalance / LAMPORTS_PER_SOL;
+
+    return {
+      sol: solBalance,
+      tokens: tokenBalances,
+    };
   }
 
   const token_account =

src/utils/tokenMetadata.ts

@@ -0,0 +1,83 @@
+import { Connection, PublicKey } from "@solana/web3.js";
+
+export async function getTokenMetadata(
+  connection: Connection,
+  tokenMint: string,
+) {
+  const METADATA_PROGRAM_ID = new PublicKey(
+    "metaqbxxUerdq28cj1RbAWkYQm3ybzjb6a8bt518x1s",
+  );
+
+  const [metadataPDA] = PublicKey.findProgramAddressSync(
+    [
+      Buffer.from("metadata"),
+      METADATA_PROGRAM_ID.toBuffer(),
+      new PublicKey(tokenMint).toBuffer(),
+    ],
+    METADATA_PROGRAM_ID,
+  );
+
+  const metadata = await connection.getAccountInfo(metadataPDA);
+  if (!metadata?.data) {
+    throw new Error("Metadata not found");
+  }
+
+  let offset = 1 + 32 + 32; // key + update auth + mint
+  const data = metadata.data;
+  const decoder = new TextDecoder();
+
+  // Read variable length strings
+  const readString = () => {
+    let nameLength = data[offset];
+
+    while (nameLength === 0) {
+      offset++;
+      nameLength = data[offset];
+      if (offset >= data.length) {
+        return null;
+      }
+    }
+
+    offset++;
+    const name = decoder
+      .decode(data.slice(offset, offset + nameLength))
+      // @eslint-disable-next-line no-control-regex
+      .replace(new RegExp(String.fromCharCode(0), "g"), "");
+    offset += nameLength;
+    return name;
+  };
+
+  const name = readString();
+  const symbol = readString();
+  const uri = readString();
+
+  // Read remaining data
+  const sellerFeeBasisPoints = data.readUInt16LE(offset);
+  offset += 2;
+
+  let creators:
+    | { address: PublicKey; verified: boolean; share: number }[]
+    | null = null;
+  if (data[offset] === 1) {
+    offset++;
+    const numCreators = data[offset];
+    offset++;
+    creators = [...Array(numCreators)].map(() => {
+      const creator = {
+        address: new PublicKey(data.slice(offset, offset + 32)),
+        verified: data[offset + 32] === 1,
+        share: data[offset + 33],
+      };
+      offset += 34;
+      return creator;
+    });
+  }
+
+  return {
+    name,
+    symbol,
+    uri,
+    sellerFeeBasisPoints,
+    creators,
+  };
+}