feat + fix: add action to get user drift account info and fix deposit, borrow, lend and withdraw on drift

src/actions/drift/depositToDriftUserAccount.ts

@@ -10,6 +10,7 @@ const depositToDriftUserAccountAction: Action = {
     "deposit into drift user account",
     "add funds to drift user account",
     "add funds to my drift account",
+    "deposit collateral into drift account",
   ],
   examples: [
     [
@@ -27,22 +28,6 @@ const depositToDriftUserAccountAction: Action = {
         explanation: "Deposit 100 USDC into your drift user account",
       },
     ],
-    [
-      {
-        input: {
-          amount: 100,
-          symbol: "USDC",
-          address: "2nFeP7taii3wGVgrWk4YiLMPmhtu3Zg9iXCUu4zGBD",
-        },
-        output: {
-          status: "success",
-          message: "Funds deposited successfully",
-          signature:
-            "2nFeP7taii3wGVgrWk4YiLMPmhtu3Zg9iXCUu4zGBDadwunHw8reXFxRWT7khbFsQ9JT3zK4RYDLNDFDRYvM3wJk",
-        },
-        explanation: "Deposit 100 USDC into a drift user account",
-      },
-    ],
   ],
   schema: z.object({
     amount: z
@@ -55,7 +40,11 @@ const depositToDriftUserAccountAction: Action = {
       .string()
       .toUpperCase()
       .describe("The symbol of the token you'd like to deposit"),
-    address: z.string().optional().describe("The drift user account address"),
+    repay: z
+      .boolean()
+      .optional()
+      .default(false)
+      .describe("Whether or not to repay the borrowed funds in the account"),
   }),
   handler: async (agent: SolanaAgentKit, input) => {
     try {
@@ -63,7 +52,7 @@ const depositToDriftUserAccountAction: Action = {
         agent,
         input.amount as number,
         input.symbol as string,
-        input.address,
+        input.repay as boolean,
       );
 
       return {

src/actions/drift/driftUserAccountInfo.ts

@@ -0,0 +1,39 @@
+import { z } from "zod";
+import type { Action } from "../../types";
+import { driftUserAccountInfo } from "../../tools";
+
+const driftUserAccountInfoAction: Action = {
+  name: "DRIFT_USER_ACCOUNT_INFO",
+  similes: ["get drift user account info", "get drift account info"],
+  description: "Get information about your drift account",
+  examples: [
+    [
+      {
+        input: {},
+        explanation: "Get information about your drift account",
+        output: {
+          status: "success",
+          data: {},
+        },
+      },
+    ],
+  ],
+  schema: z.object({}),
+  handler: async (agent) => {
+    try {
+      const accountInfo = await driftUserAccountInfo(agent);
+      return {
+        status: "success",
+        data: accountInfo,
+      };
+    } catch (e) {
+      return {
+        status: "error",
+        // @ts-expect-error - error message is a string
+        message: `Failed to get drift account info: ${e.message}`,
+      };
+    }
+  },
+};
+
+export default driftUserAccountInfoAction;

src/actions/drift/tradePerpAccount.ts

@@ -47,7 +47,10 @@ export const tradeDriftPerpAccountAction: Action = {
   ],
   schema: z.object({
     amount: z.number().positive(),
-    symbol: z.string().min(3).max(10),
+    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(),

src/actions/drift/withdrawFromDriftAccount.ts

@@ -3,12 +3,16 @@ import type { Action } from "../../types";
 import { withdrawFromDriftUserAccount } from "../../tools";
 
 const withdrawFromDriftAccountAction: Action = {
-  name: "WITHDRAW_FROM_DRIFT_ACCOUNT",
+  name: "WITHDRAW_OR_BORROW_FROM_DRIFT_ACCOUNT",
   description: "Withdraw funds from your drift account",
   similes: [
     "withdraw from drift account",
     "withdraw funds from drift account",
     "withdraw funds from my drift account",
+    "borrow from drift account",
+    "borrow funds from my drift account",
+    "borrow from drift",
+    "withdraw from drift",
   ],
   examples: [
     [
@@ -38,6 +42,13 @@ const withdrawFromDriftAccountAction: Action = {
       .string()
       .toUpperCase()
       .describe("The symbol of the token you'd like to withdraw"),
+    isBorrow: z
+      .boolean()
+      .optional()
+      .default(false)
+      .describe(
+        "Whether or not to borrow funds based on collateral provided instead of withdrawing",
+      ),
   }),
   handler: async (agent, input) => {
     try {
@@ -45,6 +56,7 @@ const withdrawFromDriftAccountAction: Action = {
         agent,
         input.amount,
         input.symbol,
+        input.isBorrow,
       );
 
       return {

src/actions/index.ts

@@ -42,6 +42,7 @@ import tradeDriftPerpAccountAction from "./drift/tradePerpAccount";
 import doesUserHaveDriftAccountAction from "./drift/doesUserHaveDriftAccount";
 import depositToDriftUserAccountAction from "./drift/depositToDriftUserAccount";
 import withdrawFromDriftAccountAction from "./drift/withdrawFromDriftAccount";
+import driftUserAccountInfoAction from "./drift/driftUserAccountInfo";
 
 export const ACTIONS = {
   WALLET_ADDRESS_ACTION: getWalletAddressAction,
@@ -88,7 +89,8 @@ export const ACTIONS = {
   TRADE_DRIFT_PERP_ACCOUNT_ACTION: tradeDriftPerpAccountAction,
   DOES_USER_HAVE_DRIFT_ACCOUNT_ACTION: doesUserHaveDriftAccountAction,
   DEPOSIT_TO_DRIFT_USER_ACCOUNT_ACTION: depositToDriftUserAccountAction,
-  WITHDRAW_FROM_DRIFT_ACCOUNT_ACTION: withdrawFromDriftAccountAction,
+  WITHDRAW_OR_BORROW_FROM_DRIFT_ACCOUNT_ACTION: withdrawFromDriftAccountAction,
+  DRIFT_USER_ACCOUNT_INFO_ACTION: driftUserAccountInfoAction,
 };
 
 export type { Action, ActionExample, Handler } from "../types/action";

src/tools/drift.ts

@@ -12,15 +12,17 @@ import {
   PositionDirection,
   PostOnlyParams,
   PRICE_PRECISION,
+  QUOTE_PRECISION,
   User,
   type IWallet,
 } from "@drift-labs/sdk";
 import type { SolanaAgentKit } from "../agent";
 import * as anchor from "@coral-xyz/anchor";
 import { IDL, VAULT_PROGRAM_ID, VaultClient } from "@drift-labs/vaults-sdk";
-import { BN } from "bn.js";
 import { getAssociatedTokenAddressSync } from "@solana/spl-token";
 import { PublicKey } from "@solana/web3.js";
+import { Transaction } from "@solana/web3.js";
+import { ComputeBudgetProgram } from "@solana/web3.js";
 
 export async function initClients(agent: SolanaAgentKit) {
   const wallet: IWallet = {
@@ -103,7 +105,7 @@ export async function createDriftUserAccount(
     }
 
     if (!userAccountExists) {
-      const depositAmount = new BN(amount).mul(token.precision);
+      const depositAmount = numberToSafeBN(amount, token.precision);
       const [txSignature, account] =
         await driftClient.initializeUserAccountAndDepositCollateral(
           depositAmount,
@@ -130,18 +132,18 @@ export async function createDriftUserAccount(
  * @param agent
  * @param amount
  * @param symbol
- * @param address
+ * @param isRepay
  * @returns
  */
 export async function depositToDriftUserAccount(
   agent: SolanaAgentKit,
   amount: number,
   symbol: string,
-  address?: string,
+  isRepay = false,
 ) {
   try {
     const { driftClient, cleanUp } = await initClients(agent);
-    const publicKey = address ? new PublicKey(address) : agent.wallet.publicKey;
+    const publicKey = agent.wallet.publicKey;
     const user = new User({
       driftClient,
       userAccountPublicKey: getUserAccountPublicKeySync(
@@ -162,11 +164,29 @@ export async function depositToDriftUserAccount(
       throw new Error("You need to create a Drift user account first.");
     }
 
-    const depositAmount = new BN(amount).mul(token.precision);
-    const txSignature = await driftClient.deposit(
-      depositAmount,
-      token.marketIndex,
-      getAssociatedTokenAddressSync(token.mint, publicKey),
+    const depositAmount = numberToSafeBN(amount, token.precision);
+
+    const [depInstruction, latestBlockhash] = await Promise.all([
+      driftClient.getDepositTxnIx(
+        depositAmount,
+        token.marketIndex,
+        getAssociatedTokenAddressSync(token.mint, publicKey),
+        undefined,
+        isRepay,
+      ),
+      driftClient.connection.getLatestBlockhash(),
+    ]);
+
+    const tx = new Transaction().add(...depInstruction).add(
+      ComputeBudgetProgram.setComputeUnitPrice({
+        microLamports: 0.000001 * 1000000 * 1000000,
+      }),
+    );
+    tx.recentBlockhash = latestBlockhash.blockhash;
+    tx.sign(agent.wallet);
+    const txSignature = await driftClient.txSender.sendRawTransaction(
+      tx.serialize(),
+      { ...driftClient.opts },
     );
 
     await cleanUp();
@@ -181,6 +201,7 @@ export async function withdrawFromDriftUserAccount(
   agent: SolanaAgentKit,
   amount: number,
   symbol: string,
+  isBorrow = false,
 ) {
   try {
     const { driftClient, cleanUp } = await initClients(agent);
@@ -207,10 +228,27 @@ export async function withdrawFromDriftUserAccount(
 
     const withdrawAmount = numberToSafeBN(amount, token.precision);
 
-    const txSignature = await driftClient.withdraw(
-      withdrawAmount,
-      token.marketIndex,
-      getAssociatedTokenAddressSync(token.mint, agent.wallet.publicKey),
+    const [withdrawInstruction, latestBlockhash] = await Promise.all([
+      driftClient.getWithdrawalIxs(
+        withdrawAmount,
+        token.marketIndex,
+        getAssociatedTokenAddressSync(token.mint, agent.wallet.publicKey),
+        !isBorrow,
+      ),
+      driftClient.connection.getLatestBlockhash(),
+    ]);
+
+    const tx = new Transaction().add(...withdrawInstruction).add(
+      ComputeBudgetProgram.setComputeUnitPrice({
+        microLamports: 0.000001 * 1000000 * 1000000,
+      }),
+    );
+    tx.recentBlockhash = latestBlockhash.blockhash;
+    tx.sign(agent.wallet);
+
+    const txSignature = await driftClient.txSender.sendRawTransaction(
+      tx.serialize(),
+      { ...driftClient.opts },
     );
 
     await cleanUp();
@@ -289,6 +327,9 @@ export async function driftPerpTrade(
           price: numberToSafeBN(params.price, PRICE_PRECISION),
           postOnly: PostOnlyParams.SLIDE,
         }),
+        {
+          computeUnitsPrice: 0.000001 * 1000000 * 1000000,
+        },
       );
     } else {
       signature = await driftClient.placePerpOrder(
@@ -301,6 +342,9 @@ export async function driftPerpTrade(
               : PositionDirection.SHORT,
           marketIndex: market.marketIndex,
         }),
+        {
+          computeUnitsPrice: 0.000001 * 1000000 * 1000000,
+        },
       );
     }
 
@@ -342,3 +386,61 @@ export async function doesUserHaveDriftAccount(agent: SolanaAgentKit) {
     throw new Error(`Failed to check user account: ${e.message}`);
   }
 }
+
+/**
+ * Get account info for a drift User
+ * @param agent
+ * @returns
+ */
+export async function driftUserAccountInfo(agent: SolanaAgentKit) {
+  try {
+    const { driftClient, cleanUp } = await initClients(agent);
+    const user = new User({
+      driftClient,
+      userAccountPublicKey: getUserAccountPublicKeySync(
+        new PublicKey(DRIFT_PROGRAM_ID),
+        agent.wallet.publicKey,
+      ),
+    });
+    const userAccountExists = await user.exists();
+
+    if (!userAccountExists) {
+      throw new Error("User account does not exist");
+    }
+    await user.subscribe();
+    const account = user.getUserAccount();
+    await user.unsubscribe();
+
+    await cleanUp();
+    const perpPositions = account.perpPositions.map((pos) => ({
+      ...pos,
+      baseAssetAmount: convertToNumber(pos.baseAssetAmount, BASE_PRECISION),
+      settledPnl: convertToNumber(pos.settledPnl, QUOTE_PRECISION),
+    }));
+    const spotPositions = account.spotPositions.map((pos) => ({
+      ...pos,
+      scaledBalance: convertToNumber(pos.scaledBalance, BASE_PRECISION),
+      cumulativeDeposits: convertToNumber(
+        pos.cumulativeDeposits,
+        BASE_PRECISION,
+      ),
+      symbol: MainnetSpotMarkets.find((v) => v.marketIndex === pos.marketIndex)
+        ?.symbol,
+    }));
+
+    return {
+      ...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,
+      spotPositions,
+    };
+  } catch (e) {
+    // @ts-expect-error - error message is a string
+    throw new Error(`Failed to check user account: ${e.message}`);
+  }
+}

src/tools/drift_vault.ts

@@ -134,7 +134,7 @@ export async function createVault(
         .toNumber(),
       minDepositAmount: numberToSafeBN(params.minDepositAmount, spotPrecision),
       redeemPeriod: new BN(params.redeemPeriod * 86400),
-      maxTokens: new BN(params.maxTokens).mul(spotPrecision),
+      maxTokens: numberToSafeBN(params.maxTokens, spotPrecision),
       managementFee: new BN(params.managementFee)
         .mul(PERCENTAGE_PRECISION)
         .div(new BN(100)),