feat: add more drift actions

src/actions/drift/availableMarkets.ts

@@ -0,0 +1,55 @@
+import { MainnetSpotMarkets } from "@drift-labs/sdk";
+import type { Action } from "../../types";
+import { z } from "zod";
+import {
+  getAvailableDriftPerpMarkets,
+  getAvailableDriftSpotMarkets,
+} from "../../tools";
+
+const availableDriftMarketsAction: Action = {
+  name: "AVAILABLE_DRIFT_MARKETS",
+  description: "Get a list of available drift markets",
+  similes: [
+    "get drift markets",
+    "drift markets",
+    "available drift markets",
+    "get available drift perp markets",
+    "get available spot markets on drift",
+  ],
+  examples: [
+    [
+      {
+        input: {
+          marketType: "spot",
+        },
+        output: {
+          status: "success",
+          message: `The list of available spot markets are ${MainnetSpotMarkets.map((v) => v.symbol).join(", ")}`,
+          data: MainnetSpotMarkets,
+        },
+        explanation: "Get the list of available spot markets/tokens on drift",
+      },
+    ],
+  ],
+  schema: z.object({
+    marketType: z
+      .enum(["spot", "perp"])
+      .describe("Type of market to get")
+      .optional(),
+  }),
+  handler: async (agent, input) => {
+    switch (input.marketType) {
+      case "perp":
+        return getAvailableDriftPerpMarkets();
+      case "spot":
+        return getAvailableDriftSpotMarkets();
+      default:
+        return {
+          spot: getAvailableDriftSpotMarkets(),
+          perp: getAvailableDriftPerpMarkets(),
+        };
+    }
+  },
+};
+
+export default availableDriftMarketsAction;

src/actions/drift/createDriftUserAccount.ts

@@ -27,7 +27,12 @@ const createDriftUserAccountAction: Action = {
     ],
   ],
   schema: z.object({
-    amount: z.number().positive().describe("Amount of the token to deposit"),
+    amount: z
+      .number()
+      .positive()
+      .describe(
+        "Amount of the token to deposit. In normal token amounts e.g 50 SOL, 100 USDC, etc",
+      ),
     symbol: z.string().describe("Symbol of the token to deposit"),
   }),
   handler: async (agent, input) => {

src/actions/drift/createVault.ts

@@ -53,9 +53,14 @@ const createDriftVaultAction: Action = {
       .int()
       .min(100, "Max tokens must be at least 100")
       .describe(
-        "The maximum amount of tokens the vault will be accomodating. For example some vaults have a cap at 10 million USDC",
+        "The maximum amount of tokens the vault will be accomodating. For example some vaults have a cap at 10 million USDC. This amount should be normal token amounts e.g 50 SOL, 100 USDC, etc",
+      ),
+    minDepositAmount: z
+      .number()
+      .positive()
+      .describe(
+        "Minimum deposit amount in normal token values e.g 50 SOL, 100 USDC, etc",
       ),
-    minDepositAmount: z.number().positive().describe("Minimum deposit amount"),
     managementFee: z
       .number()
       .positive()

src/actions/drift/depositIntoVault.ts

@@ -28,7 +28,9 @@ const depositIntoDriftVaultAction: Action = {
     amount: z
       .number()
       .positive()
-      .describe("The amount in tokens you'd like to deposit into the vault"),
+      .describe(
+        "The amount in tokens you'd like to deposit into the vault in normal token amounts e.g 50 SOL, 100 USDC, etc",
+      ),
   }),
   handler: async (agent, input) => {
     try {

src/actions/drift/depositToDriftUserAccount.ts

@@ -34,7 +34,7 @@ const depositToDriftUserAccountAction: Action = {
       .number()
       .positive()
       .describe(
-        "The amount in tokens you'd like to deposit into your drift user account",
+        "The amount in tokens you'd like to deposit into your drift user account in normal token amounts e.g 50 SOL, 100 USDC, etc",
       ),
     symbol: z
       .string()

src/actions/drift/requestUnstakeFromDriftInsuranceFund.ts

@@ -0,0 +1,62 @@
+import { z } from "zod";
+import type { Action } from "../../types";
+import { requestUnstakeFromDriftInsuranceFund } from "../../tools";
+
+const requestUnstakeFromDriftInsuranceFundAction: Action = {
+  name: "REQUEST_UNSTAKE_FROM_DRIFT_INSURANCE_FUND_ACTION",
+  description:
+    "Request to unstake a certain amount of a token from the Drift Insurance Fund",
+  similes: [
+    "request an unstake from the drift insurance fund",
+    "unstake an amount from the drift insurance fund",
+    "ask to unstake a certain amount from the drift insurance fund",
+  ],
+  examples: [
+    [
+      {
+        input: {
+          amount: 100,
+          symbol: "SOL",
+        },
+        output: {
+          status: "success",
+          message: "Requested to unstake 100 SOL from the Drift Insurance Fund",
+          signature: "4FdasklhiIHyOI",
+        },
+        explanation: "Request to unstake 100 SOL from the Drift Insurance Fund",
+      },
+    ],
+  ],
+  schema: z.object({
+    amount: z
+      .number()
+      .positive()
+      .describe("Amount to unstake in normal units e.g 50 === 50 SOL"),
+    symbol: z.string().describe("Symbol of the token to unstake"),
+  }),
+  handler: async (agent, input) => {
+    try {
+      const tx = await requestUnstakeFromDriftInsuranceFund(
+        agent,
+        input.amount,
+        input.symbol,
+      );
+
+      return {
+        status: "success",
+        message: `Requested to unstake ${input.amount} ${input.symbol} from the Drift Insurance Fund`,
+        data: {
+          signature: tx,
+        },
+      };
+    } catch (e) {
+      return {
+        status: "error",
+        // @ts-expect-error error is not a string
+        message: e.message,
+      };
+    }
+  },
+};
+
+export default requestUnstakeFromDriftInsuranceFundAction;

src/actions/drift/requestWithdrawalFromVault.ts

@@ -29,7 +29,9 @@ const requestWithdrawalFromVaultAction: Action = {
     amount: z
       .number()
       .positive()
-      .describe("Amount of shares you would like to withdraw from the vault"),
+      .describe(
+        "Amount of shares you would like to withdraw from the vault in normal token amounts e.g 50 SOL, 100 USDC, etc",
+      ),
   }),
   handler: async (agent: SolanaAgentKit, input) => {
     try {

src/actions/drift/stakeToDriftInsuranceFund.ts

@@ -0,0 +1,59 @@
+import { z } from "zod";
+import type { Action } from "../../types";
+import { stakeToDriftInsuranceFund } from "../../tools";
+
+const stakeToDriftInsuranceFundAction: Action = {
+  name: "STAKE_TO_DRIFT_INSURANCE_FUND_ACTION",
+  description: "Stake a token to Drift Insurance Fund",
+  similes: ["Stake a token to Drift Insurance Fund"],
+  examples: [
+    [
+      {
+        input: {
+          amount: 100,
+          symbol: "SOL",
+        },
+        output: {
+          status: "success",
+          message: "Staked 100 SOL to the Drift Insurance Fund",
+          data: {
+            signature: "signature",
+          },
+        },
+        explanation: "Stake 100 SOL to the Drift Insurance Fund",
+      },
+    ],
+  ],
+  schema: z.object({
+    amount: z
+      .number()
+      .positive()
+      .describe("Amount to stake in normal units e.g 50 === 50 SOL"),
+    symbol: z.string().describe("Symbol of the token stake"),
+  }),
+  handler: async (agent, input) => {
+    try {
+      const tx = await stakeToDriftInsuranceFund(
+        agent,
+        input.amount,
+        input.symbol,
+      );
+
+      return {
+        status: "sucess",
+        message: `Staked ${input.amount} ${input.symbol} to the Drift Insurance Fund`,
+        data: {
+          signature: tx,
+        },
+      };
+    } catch (error) {
+      return {
+        status: "error",
+        // @ts-expect-error error is not a string
+        message: error.message,
+      };
+    }
+  },
+};
+
+export default stakeToDriftInsuranceFundAction;

src/actions/drift/swapSpotToken.ts

@@ -0,0 +1,76 @@
+import { z } from "zod";
+import type { Action } from "../../types";
+import { swapSpotToken } from "../../tools";
+
+const driftSpotTokenSwapAction: Action = {
+  name: "DRIFT_SPOT_TOKEN_SWAP_ACTION",
+  description: "Swap a spot token for another spot token on Drift",
+  similes: [
+    "swap a token for another token on drift",
+    "exchange a token for another token on drift",
+    "trade a token for another token on drift",
+  ],
+  examples: [
+    [
+      {
+        input: {
+          fromSymbol: "SOL",
+          toSymbol: "USDC",
+          fromAmount: 100,
+        },
+        output: {
+          status: "success",
+          message: "Swapped 100 SOL for USDC on Drift",
+          signature: "4FdasklhiIHyOI",
+        },
+        explanation: "Swap 100 SOL for USDC on Drift",
+      },
+    ],
+  ],
+  schema: z.object({
+    fromSymbol: z.string().describe("Symbol of the token to swap from"),
+    toSymbol: z.string().describe("Symbol of the token to swap to"),
+    fromAmount: z
+      .number()
+      .positive()
+      .describe("Amount to swap from in normal units e.g 50 === 50 SOL")
+      .optional(),
+    toAmount: z
+      .number()
+      .positive()
+      .describe("Amount to swap to in normal units e.g 5000 === 5000 USDC")
+      .optional(),
+    slippage: z
+      .number()
+      .positive()
+      .describe("Slippage tolerance in percentage e.g 0.5 === 0.5%")
+      .optional(),
+  }),
+  handler: async (agent, input) => {
+    try {
+      const tx = await swapSpotToken(agent, {
+        fromSymbol: input.fromSymbol,
+        toSymbol: input.toSymbol,
+        fromAmount: input.fromAmount,
+        toAmount: input.toAmount,
+        slippage: input.slippage,
+      });
+
+      return {
+        status: "success",
+        message: `Swapped ${input.fromAmount} ${input.fromSymbol} for ${input.toAmount} ${input.toSymbol} on Drift`,
+        data: {
+          signature: tx,
+        },
+      };
+    } catch (e) {
+      return {
+        status: "error",
+        // @ts-expect-error error is not a string
+        message: e.message,
+      };
+    }
+  },
+};
+
+export default driftSpotTokenSwapAction;

src/actions/drift/tradeDelegatedDriftVault.ts

@@ -64,11 +64,20 @@ const tradeDelegatedDriftVaultAction: Action = {
   ],
   schema: z.object({
     vaultAddress: z.string().describe("Address of the Drift vault to trade in"),
-    amount: z.number().positive().describe("Amount to trade"),
+    amount: z
+      .number()
+      .positive()
+      .describe(
+        "Amount to trade in normal token amounts e.g 50 SOL, 100 USDC, etc",
+      ),
     symbol: z.string().describe("Symbol of the token to trade"),
     action: z.enum(["long", "short"]).describe("Trade action - long or short"),
     type: z.enum(["market", "limit"]).describe("Trade type - market or limit"),
-    price: z.number().positive().optional().describe("Price for limit order"),
+    price: z
+      .number()
+      .positive()
+      .optional()
+      .describe("USD price for limit order"),
   }),
   handler: async (agent: SolanaAgentKit, input) => {
     try {

src/actions/drift/tradePerpAccount.ts

@@ -46,14 +46,27 @@ export const tradeDriftPerpAccountAction: Action = {
     ],
   ],
   schema: z.object({
-    amount: z.number().positive(),
+    amount: z
+      .number()
+      .positive()
+      .describe(
+        "The amount of the token to trade in normal token amounts e.g 50 SOL, 100 USDC",
+      ),
     symbol: z
       .string()
       .toUpperCase()
       .describe("Symbol of the token to open a position on "),
-    action: z.enum(["long", "short"]),
-    type: z.enum(["market", "limit"]),
-    price: z.number().positive().optional(),
+    action: z
+      .enum(["long", "short"])
+      .describe(
+        "The action you would like to carry out whether it be a long or a short",
+      ),
+    type: z
+      .enum(["market", "limit"])
+      .describe(
+        "The type of trade you would like to open, market or limit order",
+      ),
+    price: z.number().positive().optional().describe("USD price of the token"),
   }),
   handler: async (agent, input) => {
     try {

src/actions/drift/unstakeFromDriftInsuranceFund.ts

@@ -0,0 +1,51 @@
+import { z } from "zod";
+import type { Action } from "../../types";
+import { unstakeFromDriftInsuranceFund } from "../../tools";
+
+const unstakeFromDriftInsuranceFundAction: Action = {
+  name: "UNSTAKE_FROM_DRIFT_INSURANCE_FUND_ACTION",
+  description:
+    "Unstake requested unstake amount from the Drift Insurance fund once the cool period has elapsed",
+  similes: [
+    "unstake from the drift insurance fund",
+    "withdraw from the drift insurance fund",
+    "take out funds from the drift insurance fund",
+  ],
+  examples: [
+    [
+      {
+        input: {
+          symbol: "SOL",
+        },
+        output: {
+          status: "success",
+          message: "Unstaked your SOL from the Drift Insurance Fund",
+          signature: "4FdasklhiIHyOI",
+        },
+        explanation: "Unstake SOL from the Drift Insurance Fund",
+      },
+    ],
+  ],
+  schema: z.object({
+    symbol: z.string().describe("Symbol of the token to unstake"),
+  }),
+  handler: async (agent, input) => {
+    try {
+      const tx = await unstakeFromDriftInsuranceFund(agent, input.symbol);
+
+      return {
+        status: "success",
+        message: `Unstaked your ${input.symbol} from the Drift Insurance Fund`,
+        signature: tx,
+      };
+    } catch (e) {
+      return {
+        status: "error",
+        // @ts-expect-error error is not a string
+        message: e.message,
+      };
+    }
+  },
+};
+
+export default unstakeFromDriftInsuranceFundAction;

src/actions/drift/updateDriftVaultDelegate.ts

@@ -24,8 +24,8 @@ const updateDriftVaultDelegateAction: Action = {
     ],
   ],
   schema: z.object({
-    vaultAddress: z.string(),
-    newDelegate: z.string(),
+    vaultAddress: z.string().describe("vault's address"),
+    newDelegate: z.string().describe("new address to delegate the vault to"),
   }),
   handler: async (agent, input) => {
     try {

src/actions/drift/updateVault.ts

@@ -41,16 +41,35 @@ const updateDriftVaultAction: Action = {
     redeemPeriod: z
       .number()
       .int()
-      .min(1, "Redeem period must be at least 1")
+      .min(1, "Redeem period must be at least 1 day")
       .optional(),
     maxTokens: z
       .number()
       .int()
-      .min(100, "Max tokens must be at least 100")
-      .optional(),
-    minDepositAmount: z.number().positive().optional(),
-    managementFee: z.number().positive().max(20).optional(),
-    profitShare: z.number().positive().max(90).optional(),
+      .min(100, "Max tokens must be at least be 100 units")
+      .optional()
+      .describe(
+        "The maximum number of tokens the vault is willing to accept and manage",
+      ),
+    minDepositAmount: z
+      .number()
+      .positive()
+      .optional()
+      .describe(
+        "The minimum amount that is allowed to be deposited into the vault in normal token amounts e.g 10 USDC",
+      ),
+    managementFee: z
+      .number()
+      .positive()
+      .max(20)
+      .optional()
+      .describe("The percentage fee the vault takes for asset management"),
+    profitShare: z
+      .number()
+      .positive()
+      .max(90)
+      .optional()
+      .describe("Profit share in percentage e.g 2 === 2%"),
     handleRate: z.number().optional(),
     permissioned: z
       .boolean()

src/actions/drift/withdrawFromDriftAccount.ts

@@ -36,7 +36,7 @@ const withdrawFromDriftAccountAction: Action = {
       .number()
       .positive()
       .describe(
-        "The amount in tokens you'd like to withdraw from your drift account",
+        "The amount in tokens you'd like to withdraw from your drift account in normal token amounts, e.g 50 SOL, 100 USDC, etc",
       ),
     symbol: z
       .string()

src/actions/drift/withdrawFromVault.ts

@@ -25,7 +25,7 @@ const withdrawFromVaultAction: Action = {
     ],
   ],
   schema: z.object({
-    vaultAddress: z.string(),
+    vaultAddress: z.string().describe("Vault's address"),
   }),
   handler: async (agent: SolanaAgentKit, input) => {
     try {

src/actions/index.ts

@@ -59,6 +59,11 @@ import withdrawFromDriftAccountAction from "./drift/withdrawFromDriftAccount";
 import driftUserAccountInfoAction from "./drift/driftUserAccountInfo";
 import deriveDriftVaultAddressAction from "./drift/deriveVaultAddress";
 import updateDriftVaultDelegateAction from "./drift/updateDriftVaultDelegate";
+import availableDriftMarketsAction from "./drift/availableMarkets";
+import stakeToDriftInsuranceFundAction from "./drift/stakeToDriftInsuranceFund";
+import requestUnstakeFromDriftInsuranceFundAction from "./drift/requestUnstakeFromDriftInsuranceFund";
+import unstakeFromDriftInsuranceFundAction from "./drift/unstakeFromDriftInsuranceFund";
+import driftSpotTokenSwapAction from "./drift/swapSpotToken";
 
 export const ACTIONS = {
   WALLET_ADDRESS_ACTION: getWalletAddressAction,
@@ -123,6 +128,12 @@ export const ACTIONS = {
   DRIFT_USER_ACCOUNT_INFO_ACTION: driftUserAccountInfoAction,
   DERIVE_DRIFT_VAULT_ADDRESS_ACTION: deriveDriftVaultAddressAction,
   UPDATE_DRIFT_VAULT_DELEGATE_ACTION: updateDriftVaultDelegateAction,
+  AVAILABLE_DRIFT_MARKETS_ACTION: availableDriftMarketsAction,
+  STAKE_TO_DRIFT_INSURANCE_FUND_ACTION: stakeToDriftInsuranceFundAction,
+  REQUEST_UNSTAKE_FROM_DRIFT_INSURANCE_FUND_ACTION:
+    requestUnstakeFromDriftInsuranceFundAction,
+  UNSTAKE_FROM_DRIFT_INSURANCE_FUND_ACTION: unstakeFromDriftInsuranceFundAction,
+  DRIFT_SPOT_TOKEN_SWAP_ACTION: driftSpotTokenSwapAction,
 };
 
 export type { Action, ActionExample, Handler } from "../types/action";

src/agent/index.ts

@@ -98,6 +98,12 @@ import {
   withdrawFromDriftVault,
   updateVaultDelegate,
   get_token_balance,
+  getAvailableDriftSpotMarkets,
+  getAvailableDriftPerpMarkets,
+  stakeToDriftInsuranceFund,
+  requestUnstakeFromDriftInsuranceFund,
+  unstakeFromDriftInsuranceFund,
+  swapSpotToken,
 } from "../tools";
 import {
   Config,
@@ -821,4 +827,48 @@ export class SolanaAgentKit {
   async updateDriftVaultDelegate(vaultAddress: string, delegate: string) {
     return await updateVaultDelegate(this, vaultAddress, delegate);
   }
+  getAvailableDriftMarkets(type?: "spot" | "perp") {
+    switch (type) {
+      case "spot":
+        return getAvailableDriftSpotMarkets();
+      case "perp":
+        return getAvailableDriftPerpMarkets();
+      default:
+        return {
+          spot: getAvailableDriftSpotMarkets(),
+          perp: getAvailableDriftPerpMarkets(),
+        };
+    }
+  }
+  async stakeToDriftInsuranceFund(amount: number, symbol: string) {
+    return await stakeToDriftInsuranceFund(this, amount, symbol);
+  }
+  async requestUnstakeFromDriftInsuranceFund(amount: number, symbol: string) {
+    return await requestUnstakeFromDriftInsuranceFund(this, amount, symbol);
+  }
+  async unstakeFromDriftInsuranceFund(symbol: string) {
+    return await unstakeFromDriftInsuranceFund(this, symbol);
+  }
+  async driftSpotTokenSwap(
+    params: {
+      fromSymbol: string;
+      toSymbol: string;
+      slippage?: number;
+    } & (
+      | {
+          toAmount: number;
+        }
+      | { fromAmount: number }
+    ),
+  ) {
+    return await swapSpotToken(this, {
+      fromSymbol: params.fromSymbol,
+      toSymbol: params.toSymbol,
+      // @ts-expect-error - fromAmount and toAmount are mutually exclusive
+      fromAmount: params.fromAmount,
+      // @ts-expect-error - fromAmount and toAmount are mutually exclusive
+      toAmount: params.toAmount,
+      slippage: params.slippage,
+    });
+  }
 }

src/langchain/drift/request_unstake_from_insurance_fund.ts

@@ -0,0 +1,37 @@
+import { Tool } from "langchain/tools";
+import type { SolanaAgentKit } from "../../agent";
+
+export class SolanaRequestUnstakeFromDriftInsuranceFundTool extends Tool {
+  name = "request_unstake_from_drift_insurance_fund";
+  description = `Request to unstake tokens from Drift Insurance Fund.
+  
+  Inputs (JSON string):
+  - amount: number, amount to unstake (required)
+  - symbol: string, token symbol (required)`;
+
+  constructor(private solanaKit: SolanaAgentKit) {
+    super();
+  }
+
+  protected async _call(input: string): Promise<string> {
+    try {
+      const parsedInput = JSON.parse(input);
+      const tx = await this.solanaKit.requestUnstakeFromDriftInsuranceFund(
+        parsedInput.amount,
+        parsedInput.symbol,
+      );
+
+      return JSON.stringify({
+        status: "success",
+        message: `Requested unstake of ${parsedInput.amount} ${parsedInput.symbol} from the Drift Insurance Fund`,
+        signature: tx,
+      });
+    } catch (error: any) {
+      return JSON.stringify({
+        status: "error",
+        message: error.message,
+        code: error.code || "REQUEST_UNSTAKE_FROM_DRIFT_INSURANCE_FUND_ERROR",
+      });
+    }
+  }
+}

src/langchain/drift/stake_to_insurance_fund.ts

@@ -0,0 +1,37 @@
+import { Tool } from "langchain/tools";
+import type { SolanaAgentKit } from "../../agent";
+
+export class SolanaStakeToDriftInsuranceFundTool extends Tool {
+  name = "stake_to_drift_insurance_fund";
+  description = `Stake a token to Drift Insurance Fund.
+  
+  Inputs (JSON string):
+  - amount: number, amount to stake (required)
+  - symbol: string, token symbol (required)`;
+
+  constructor(private solanaKit: SolanaAgentKit) {
+    super();
+  }
+
+  protected async _call(input: string): Promise<string> {
+    try {
+      const parsedInput = JSON.parse(input);
+      const tx = await this.solanaKit.stakeToDriftInsuranceFund(
+        parsedInput.amount,
+        parsedInput.symbol,
+      );
+
+      return JSON.stringify({
+        status: "success",
+        message: `Staked ${parsedInput.amount} ${parsedInput.symbol} to the Drift Insurance Fund`,
+        signature: tx,
+      });
+    } catch (error: any) {
+      return JSON.stringify({
+        status: "error",
+        message: error.message,
+        code: error.code || "STAKE_TO_DRIFT_INSURANCE_FUND_ERROR",
+      });
+    }
+  }
+}

src/langchain/drift/swap_spot_token.ts

@@ -0,0 +1,37 @@
+import { Tool } from "langchain/tools";
+import type { SolanaAgentKit } from "../../agent";
+
+export class SolanaDriftSpotTokenSwapTool extends Tool {
+  name = "drift_spot_token_swap";
+  description = `Swap spot tokens on Drift protocol.
+  
+  Inputs (JSON string):
+  - fromSymbol: string, symbol of token to swap from (required)
+  - toSymbol: string, symbol of token to swap to (required)
+  - fromAmount: number, amount to swap from (optional) required if toAmount is not provided
+  - toAmount: number, amount to swap to (optional) required if fromAmount is not provided
+  - slippage: number, slippage tolerance in percentage (optional)`;
+
+  constructor(private solanaKit: SolanaAgentKit) {
+    super();
+  }
+
+  protected async _call(input: string): Promise<string> {
+    try {
+      const parsedInput = JSON.parse(input);
+      const tx = await this.solanaKit.driftSpotTokenSwap(parsedInput);
+
+      return JSON.stringify({
+        status: "success",
+        message: `Swapped ${parsedInput.fromAmount} ${parsedInput.fromSymbol} for ${parsedInput.toSymbol}`,
+        signature: tx,
+      });
+    } catch (error: any) {
+      return JSON.stringify({
+        status: "error",
+        message: error.message,
+        code: error.code || "DRIFT_SPOT_TOKEN_SWAP_ERROR",
+      });
+    }
+  }
+}

src/langchain/drift/unstake_from_insurance_fund.ts

@@ -0,0 +1,32 @@
+import { Tool } from "langchain/tools";
+import type { SolanaAgentKit } from "../../agent";
+
+export class SolanaUnstakeFromDriftInsuranceFundTool extends Tool {
+  name = "unstake_from_drift_insurance_fund";
+  description = `Unstake tokens from Drift Insurance Fund after request period has elapsed.
+  
+  Inputs (JSON string):
+  - symbol: string, token symbol (required)`;
+
+  constructor(private solanaKit: SolanaAgentKit) {
+    super();
+  }
+
+  protected async _call(input: string): Promise<string> {
+    try {
+      const tx = await this.solanaKit.unstakeFromDriftInsuranceFund(input);
+
+      return JSON.stringify({
+        status: "success",
+        message: `Unstaked ${input} from the Drift Insurance Fund`,
+        signature: tx,
+      });
+    } catch (error: any) {
+      return JSON.stringify({
+        status: "error",
+        message: error.message,
+        code: error.code || "UNSTAKE_FROM_DRIFT_INSURANCE_FUND_ERROR",
+      });
+    }
+  }
+}

src/tools/drift/drift.ts

@@ -7,6 +7,8 @@ import {
   getLimitOrderParams,
   getMarketOrderParams,
   getUserAccountPublicKeySync,
+  JupiterClient,
+  MainnetPerpMarkets,
   MainnetSpotMarkets,
   numberToSafeBN,
   PositionDirection,
@@ -115,7 +117,10 @@ export async function createDriftUserAccount(
     );
 
     if (!token) {
-      throw new Error(`Token with symbol ${symbol} not found`);
+      throw new Error(`Token with symbol ${symbol} not found. Here's a list of available spot markets: ${MainnetSpotMarkets.map(
+        (v) => v.symbol,
+      ).join(", ")}
+      `);
     }
 
     if (!userAccountExists) {
@@ -171,7 +176,11 @@ export async function depositToDriftUserAccount(
     );
 
     if (!token) {
-      throw new Error(`Token with symbol ${symbol} not found`);
+      throw new Error(
+        `Token with symbol ${symbol} not found. Here's a list of available spot markets: ${MainnetSpotMarkets.map(
+          (v) => v.symbol,
+        ).join(", ")}`,
+      );
     }
 
     if (!userAccountExists) {
@@ -237,7 +246,11 @@ export async function withdrawFromDriftUserAccount(
     );
 
     if (!token) {
-      throw new Error(`Token with symbol ${symbol} not found`);
+      throw new Error(
+        `Token with symbol ${symbol} not found. Here's a list of available spot markets: ${MainnetSpotMarkets.map(
+          (v) => v.symbol,
+        ).join(", ")}`,
+      );
     }
 
     const withdrawAmount = numberToSafeBN(amount, token.precision);
@@ -313,7 +326,11 @@ export async function driftPerpTrade(
     );
 
     if (!market) {
-      throw new Error(`Token with symbol ${params.symbol} not found`);
+      throw new Error(
+        `Token with symbol ${params.symbol} not found. Here's a list of available perp markets: ${MainnetPerpMarkets.map(
+          (v) => v.symbol,
+        ).join(", ")}`,
+      );
     }
 
     const baseAssetPrice = driftClient.getOracleDataForPerpMarket(
@@ -435,10 +452,9 @@ export async function driftUserAccountInfo(agent: SolanaAgentKit) {
     }));
     const spotPositions = account.spotPositions.map((pos) => ({
       ...pos,
-      scaledBalance: convertToNumber(pos.scaledBalance, BASE_PRECISION),
-      cumulativeDeposits: convertToNumber(
-        pos.cumulativeDeposits,
-        BASE_PRECISION,
+      availableBalance: convertToNumber(
+        pos.scaledBalance,
+        MainnetSpotMarkets[pos.marketIndex].precision,
       ),
       symbol: MainnetSpotMarkets.find((v) => v.marketIndex === pos.marketIndex)
         ?.symbol,
@@ -448,8 +464,6 @@ export async function driftUserAccountInfo(agent: SolanaAgentKit) {
       ...account,
       name: account.name,
       authority: account.authority,
-      totalDeposits: `$${convertToNumber(account.totalDeposits, QUOTE_PRECISION)}`,
-      totalWithdraws: `$${convertToNumber(account.totalWithdraws, QUOTE_PRECISION)}`,
       settledPerpPnl: `$${convertToNumber(account.settledPerpPnl, QUOTE_PRECISION)}`,
       lastActiveSlot: account.lastActiveSlot.toNumber(),
       perpPositions,
@@ -460,3 +474,239 @@ export async function driftUserAccountInfo(agent: SolanaAgentKit) {
     throw new Error(`Failed to check user account: ${e.message}`);
   }
 }
+
+/**
+ * Get available spot markets on drift protocol
+ */
+export function getAvailableDriftSpotMarkets() {
+  return MainnetSpotMarkets;
+}
+
+/**
+ * Get available perp markets on drift protocol
+ */
+export function getAvailableDriftPerpMarkets() {
+  return MainnetPerpMarkets;
+}
+
+/**
+ * Stake a token to the drift insurance fund
+ * @param agent
+ * @param amount
+ * @param symbol
+ */
+export async function stakeToDriftInsuranceFund(
+  agent: SolanaAgentKit,
+  amount: number,
+  symbol: string,
+) {
+  try {
+    const { cleanUp, driftClient } = await initClients(agent);
+    const token = MainnetSpotMarkets.find(
+      (v) => v.symbol === symbol.toUpperCase(),
+    );
+
+    if (!token) {
+      throw new Error(
+        `Token with symbol ${symbol} not found. Here's a list of available spot markets: ${MainnetSpotMarkets.map(
+          (v) => v.symbol,
+        ).join(", ")}`,
+      );
+    }
+
+    const signature = await driftClient.addInsuranceFundStake({
+      amount: numberToSafeBN(amount, token.precision),
+      marketIndex: token.marketIndex,
+      collateralAccountPublicKey: getAssociatedTokenAddressSync(
+        token.mint,
+        agent.wallet.publicKey,
+      ),
+      txParams: {
+        computeUnitsPrice: 0.000002 * 1000000 * 1000000,
+      },
+    });
+
+    await cleanUp();
+    return signature;
+  } catch (e) {
+    // @ts-expect-error - error message is a string
+    throw new Error(`Failed to get APYs: ${e.message}`);
+  }
+}
+
+/**
+ * Request an unstake from the drift insurance fund
+ * @param agent
+ * @param amount
+ * @param symbol
+ */
+export async function requestUnstakeFromDriftInsuranceFund(
+  agent: SolanaAgentKit,
+  amount: number,
+  symbol: string,
+) {
+  try {
+    const { driftClient, cleanUp } = await initClients(agent);
+    const token = MainnetSpotMarkets.find(
+      (v) => v.symbol === symbol.toUpperCase(),
+    );
+
+    if (!token) {
+      throw new Error(
+        `Token with symbol ${symbol} not found. Here's a list of available spot markets: ${MainnetSpotMarkets.map(
+          (v) => v.symbol,
+        ).join(", ")}`,
+      );
+    }
+
+    const signature = await driftClient.requestRemoveInsuranceFundStake(
+      token.marketIndex,
+      numberToSafeBN(amount, token.precision),
+      { computeUnitsPrice: 0.000002 * 1000000 * 1000000 },
+    );
+
+    await cleanUp();
+    return signature;
+  } catch (e) {
+    // @ts-expect-error error message is a string
+    throw new Error(`Failed to unstake from insurance fund: ${e.message}`);
+  }
+}
+
+/**
+ * Unstake requested funds from the drift insurance fund once cool down period is elapsed
+ * @param agent
+ * @param symbol
+ */
+export async function unstakeFromDriftInsuranceFund(
+  agent: SolanaAgentKit,
+  symbol: string,
+) {
+  try {
+    const { driftClient, cleanUp } = await initClients(agent);
+    const token = MainnetSpotMarkets.find(
+      (v) => v.symbol === symbol.toUpperCase(),
+    );
+
+    if (!token) {
+      throw new Error(
+        `Token with symbol ${symbol} not found. Here's a list of available spot markets: ${MainnetSpotMarkets.map(
+          (v) => v.symbol,
+        ).join(", ")}`,
+      );
+    }
+
+    const signature = await driftClient.removeInsuranceFundStake(
+      token.marketIndex,
+      getAssociatedTokenAddressSync(token.mint, agent.wallet.publicKey),
+      {
+        computeUnitsPrice: 0.000002 * 1000000 * 1000000,
+      },
+    );
+
+    await cleanUp();
+    return signature;
+  } catch (e) {
+    // @ts-expect-error error message is a string
+    throw new Error(`Failed to unstake from insurance fund: ${e.message}`);
+  }
+}
+
+/**
+ * Swap a spot token for another on drift
+ * @param agent
+ * @param params
+ * @param params.fromSymbol symbol of the token to deposit
+ * @param params.toSymbol symbol of the token to receive
+ * @param params.fromAmount amount of the token to deposit
+ * @param params.toAmount amount of the token to receive
+ */
+export async function swapSpotToken(
+  agent: SolanaAgentKit,
+  params: {
+    fromSymbol: string;
+    toSymbol: string;
+    slippage?: number | undefined;
+  } & (
+    | {
+        fromAmount: number;
+      }
+    | {
+        toAmount: number;
+      }
+  ),
+) {
+  try {
+    const { driftClient, cleanUp } = await initClients(agent);
+    const fromToken = MainnetSpotMarkets.find(
+      (v) => v.symbol === params.fromSymbol.toUpperCase(),
+    );
+    const toToken = MainnetSpotMarkets.find(
+      (v) => v.symbol === params.toSymbol.toUpperCase(),
+    );
+
+    if (!fromToken) {
+      throw new Error(
+        `Token with symbol ${params.fromSymbol} not found. Here's a list of available spot markets: ${MainnetSpotMarkets.map(
+          (v) => v.symbol,
+        ).join(", ")}`,
+      );
+    }
+
+    if (!toToken) {
+      throw new Error(
+        `Token with symbol ${params.toSymbol} not found. Here's a list of available spot markets: ${MainnetSpotMarkets.map(
+          (v) => v.symbol,
+        ).join(", ")}`,
+      );
+    }
+
+    let txSig: string;
+
+    // @ts-expect-error - false undefined type conflict
+    if (params.fromAmount) {
+      const jupiterClient = new JupiterClient({ connection: agent.connection });
+      // @ts-expect-error - false undefined type conflict
+      const fromAmount = numberToSafeBN(params.fromAmount, fromToken.precision);
+      const signature = await driftClient.swap({
+        amount: fromAmount,
+        inMarketIndex: fromToken.marketIndex,
+        outMarketIndex: toToken.marketIndex,
+        jupiterClient: jupiterClient,
+        slippageBps: params.slippage ?? 100,
+        swapMode: "ExactIn",
+      });
+
+      txSig = signature;
+    }
+
+    // @ts-expect-error - false undefined type conflict
+    if (params.toAmount) {
+      const jupiterClient = new JupiterClient({ connection: agent.connection });
+      // @ts-expect-error - false undefined type conflict
+      const toAmount = numberToSafeBN(params.toAmount, toToken.precision);
+      const signature = await driftClient.swap({
+        amount: toAmount,
+        inMarketIndex: toToken.marketIndex,
+        outMarketIndex: fromToken.marketIndex,
+        jupiterClient: jupiterClient,
+        slippageBps: params.slippage ?? 100,
+        swapMode: "ExactOut",
+      });
+
+      txSig = signature;
+    }
+
+    await cleanUp();
+
+    // @ts-expect-error - false use before assignment
+    if (txSig) {
+      return txSig;
+    }
+
+    throw new Error("Either fromAmount or toAmount must be provided");
+  } catch (e) {
+    // @ts-expect-error error message is a string
+    throw new Error(`Failed to swap token: ${e.message}`);
+  }
+}

src/tools/drift/drift_vault.ts

@@ -37,14 +37,18 @@ export function getMarketIndexAndType(name: `${string}-${string}`) {
   if (type === "PERP") {
     const token = MainnetPerpMarkets.find((v) => v.baseAssetSymbol === symbol);
     if (!token) {
-      throw new Error("Drift doesn't have that market");
+      throw new Error(
+        `Drift doesn't have that market. Here's a list of available perp markets: ${MainnetPerpMarkets.map((v) => v.baseAssetSymbol).join(", ")}`,
+      );
     }
     return { marketIndex: token.marketIndex, marketType: MarketType.PERP };
   }
 
   const token = MainnetSpotMarkets.find((v) => v.symbol === symbol);
   if (!token) {
-    throw new Error("Drift doesn't have that market");
+    throw new Error(
+      `Drift doesn't have that market. Here's a list of available spot markets: ${MainnetSpotMarkets.map((v) => v.symbol).join(", ")}`,
+    );
   }
   return { marketIndex: token.marketIndex, marketType: MarketType.SPOT };
 }
@@ -134,22 +138,22 @@ export async function createVault(
     const { vaultClient, driftClient, cleanUp } = await initClients(agent);
     const marketIndexAndType = getMarketIndexAndType(params.marketName);
 
-    if (!marketIndexAndType) {
-      throw new Error("Invalid market name");
-    }
-
     const spotMarket = driftClient.getSpotMarketAccount(
       marketIndexAndType.marketIndex,
     );
 
     if (!spotMarket) {
-      throw new Error("Market not found");
+      throw new Error(
+        `Market not found. Here's a list of available spot markets: ${MainnetSpotMarkets.map((v) => `${v.symbol}-SPOT`).join(", ")}`,
+      );
     }
 
     const spotPrecision = TEN.pow(new BN(spotMarket.decimals));
 
     if (marketIndexAndType.marketType === MarketType.PERP) {
-      throw new Error("Only SPOT market names are supported");
+      throw new Error(
+        `Only SPOT market names are supported. Such as ${MainnetSpotMarkets.map((v) => `${v.symbol}-SPOT`).join(", ")}`,
+      );
     }
 
     const tx = await vaultClient.initializeVault({
@@ -239,7 +243,9 @@ export async function updateVault(
     );
 
     if (!spotMarket) {
-      throw new Error("Market not found");
+      throw new Error(
+        "Market not found. This vault's market is no longer supported",
+      );
     }
 
     const spotPrecision = TEN.pow(new BN(spotMarket.decimals));
@@ -370,7 +376,9 @@ export async function depositIntoVault(
     );
 
     if (!spotMarket) {
-      throw new Error("Market not found");
+      throw new Error(
+        "Market not found. This vaults market is no longer supported",
+      );
     }
 
     const spotPrecision = TEN.pow(new BN(spotMarket.decimals));
@@ -544,7 +552,7 @@ export async function tradeDriftVault(
 
     if (!isOwned) {
       throw new Error(
-        "This vault is owned by someone else, so you can't trade with it",
+        "This vault is owned/delegated to someone else, you can't trade with it",
       );
     }