feat: drift perp trade, account creation, withdraw, and deposit

src/actions/drift/createDriftUserAccount.ts

@@ -0,0 +1,59 @@
+import { z } from "zod";
+import type { Action } from "../../types";
+import { createDriftUserAccount } from "../../tools";
+
+const createDriftUserAccountAction: Action = {
+  name: "CREATE_DRIFT_USER_ACCOUNT",
+  similes: [
+    "create drift account",
+    "create drift user account",
+    "create user account on drift",
+  ],
+  description: "Create a new user account on Drift protocol",
+  examples: [
+    [
+      {
+        input: {
+          amount: 100,
+          symbol: "SOL",
+        },
+        output: {
+          status: "success",
+          message: "User account created with 100 SOL successfully deposited",
+          account: "4xKpN2...",
+        },
+        explanation: "Create a new user account with 100 SOL",
+      },
+    ],
+  ],
+  schema: z.object({
+    amount: z.number().positive().describe("Amount of the token to deposit"),
+    symbol: z.string().describe("Symbol of the token to deposit"),
+  }),
+  handler: async (agent, input) => {
+    try {
+      const res = await createDriftUserAccount(
+        agent,
+        input.amount,
+        input.symbol,
+      );
+
+      return {
+        status: "success",
+        message:
+          res.message ??
+          `User account created with ${input.amount} ${input.symobl} successfully deposited.`,
+        account: res.account,
+        signature: res.txSignature,
+      };
+    } catch (e) {
+      return {
+        status: "error",
+        // @ts-expect-error - error message is a string
+        message: `Failed to create user account: ${e.message}`,
+      };
+    }
+  },
+};
+
+export default createDriftUserAccountAction;

src/actions/drift/depositIntoVault.ts

@@ -1,6 +1,5 @@
 import { z } from "zod";
 import type { Action } from "../../types";
-import type { SolanaAgentKit } from "../../agent";
 import { depositIntoVault } from "../../tools/drift_vault";
 
 const depositIntoDriftVaultAction: Action = {
@@ -31,7 +30,7 @@ const depositIntoDriftVaultAction: Action = {
       .positive()
       .describe("The amount in tokens you'd like to deposit into the vault"),
   }),
-  handler: async (agent: SolanaAgentKit, input) => {
+  handler: async (agent, input) => {
     try {
       const tx = await depositIntoVault(
         agent,

src/actions/drift/depositToDriftUserAccount.ts

@@ -0,0 +1,84 @@
+import { z } from "zod";
+import type { SolanaAgentKit } from "../../agent";
+import type { Action } from "../../types";
+import { depositToDriftUserAccount } from "../../tools";
+
+const depositToDriftUserAccountAction: Action = {
+  name: "DEPOSIT_TO_DRIFT_USER_ACCOUNT",
+  description: "Deposit funds into your drift user account",
+  similes: [
+    "deposit into drift user account",
+    "add funds to drift user account",
+    "add funds to my drift account",
+  ],
+  examples: [
+    [
+      {
+        input: {
+          amount: 100,
+          symbol: "usdc",
+        },
+        output: {
+          status: "success",
+          message: "Funds deposited successfully",
+          signature:
+            "2nFeP7taii3wGVgrWk4YiLMPmhtu3Zg9iXCUu4zGBDadwunHw8reXFxRWT7khbFsQ9JT3zK4RYDLNDFDRYvM3wJk",
+        },
+        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
+      .number()
+      .positive()
+      .describe(
+        "The amount in tokens you'd like to deposit into your drift user account",
+      ),
+    symbol: z
+      .string()
+      .toUpperCase()
+      .describe("The symbol of the token you'd like to deposit"),
+    address: z.string().optional().describe("The drift user account address"),
+  }),
+  handler: async (agent: SolanaAgentKit, input) => {
+    try {
+      const tx = await depositToDriftUserAccount(
+        agent,
+        input.amount as number,
+        input.symbol as string,
+        input.address,
+      );
+
+      return {
+        status: "success",
+        message: "Funds deposited successfully",
+        signature: tx,
+      };
+    } catch (e) {
+      return {
+        status: "error",
+        // @ts-expect-error - error message
+        message: `Failed to deposit funds: ${e.message}`,
+      };
+    }
+  },
+};
+
+export default depositToDriftUserAccountAction;

src/actions/drift/doesUserHaveDriftAccount.ts

@@ -0,0 +1,53 @@
+import { z } from "zod";
+import { doesUserHaveDriftAccount } from "../../tools";
+import type { Action } from "../../types";
+
+export const doesUserHaveDriftAccountAction: Action = {
+  name: "DOES_USER_HAVE_DRIFT_ACCOUNT",
+  description: "Check if a user has a Drift account",
+  similes: [
+    "check if user has drift account",
+    "check if user has account on drift",
+    "do I have an account on drift",
+  ],
+  examples: [
+    [
+      {
+        input: {},
+        output: {
+          status: "success",
+          message: "Nice! You have a Drift account",
+          account: "4xKpN2...",
+        },
+        explanation: "Check if a user has a Drift account",
+      },
+    ],
+  ],
+  schema: z.object({}),
+  handler: async (agent) => {
+    try {
+      const res = await doesUserHaveDriftAccount(agent);
+
+      if (!res.hasAccount) {
+        return {
+          status: "error",
+          message: "You do not have a Drift account",
+        };
+      }
+
+      return {
+        status: "success",
+        message: "Nice! You have a Drift account",
+        account: res.account,
+      };
+    } catch (e) {
+      return {
+        status: "error",
+        // @ts-expect-error - error message is a string
+        message: `Failed to check if you have a Drift account: ${e.message}`,
+      };
+    }
+  },
+};
+
+export default doesUserHaveDriftAccountAction;

src/actions/drift/tradePerpAccount.ts

@@ -0,0 +1,79 @@
+import { z } from "zod";
+import type { Action } from "../../types";
+import { driftPerpTrade } from "../../tools";
+
+export const tradeDriftPerpAccountAction: Action = {
+  name: "TRADE_DRIFT_PERP_ACCOUNT",
+  similes: [
+    "trade drift perp account",
+    "trade drift perp",
+    "trade drift perpetual account",
+    "trade perp account",
+    "trade account",
+  ],
+  description: "Trade a perpetual account on Drift protocol",
+  examples: [
+    [
+      {
+        input: {
+          amount: 100,
+          symbol: "SOL",
+          action: "long",
+          type: "market",
+        },
+        output: {
+          status: "success",
+          message: "Trade successful",
+        },
+        explanation: "Open a $100 long position on SOL.",
+      },
+    ],
+    [
+      {
+        input: {
+          amount: 50,
+          symbol: "BTC",
+          action: "short",
+          type: "limit",
+          price: 50000,
+        },
+        output: {
+          status: "success",
+          message: "Trade successful",
+        },
+        explanation: "$50 short position on BTC at $50,000.",
+      },
+    ],
+  ],
+  schema: z.object({
+    amount: z.number().positive(),
+    symbol: z.string().min(3).max(10),
+    action: z.enum(["long", "short"]),
+    type: z.enum(["market", "limit"]),
+    price: z.number().positive().optional(),
+  }),
+  handler: async (agent, input) => {
+    try {
+      const signature = await driftPerpTrade(agent, {
+        action: input.action,
+        amount: input.amount,
+        symbol: input.symbol,
+        type: input.type,
+        price: input.price,
+      });
+
+      return {
+        status: "success",
+        signature: signature,
+      };
+    } catch (e) {
+      return {
+        status: "error",
+        // @ts-expect-error - error message is a string
+        message: `Failed to trade perp account: ${e.message}`,
+      };
+    }
+  },
+};
+
+export default tradeDriftPerpAccountAction;

src/actions/drift/withdrawFromDriftAccount.ts

@@ -0,0 +1,65 @@
+import { z } from "zod";
+import type { Action } from "../../types";
+import { withdrawFromDriftUserAccount } from "../../tools";
+
+const withdrawFromDriftAccountAction: Action = {
+  name: "WITHDRAW_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",
+  ],
+  examples: [
+    [
+      {
+        input: {
+          amount: 100,
+          symbol: "usdc",
+        },
+        output: {
+          status: "success",
+          message: "Funds withdrawn successfully",
+          signature:
+            "2nFeP7taii3wGVgrWk4YiLMPmhtu3Zg9iXCUu4zGBDadwunHw8reXFxRWT7khbFsQ9JT3zK4RYDLNDFDRYvM3wJk",
+        },
+        explanation: "Withdraw 100 USDC from your drift account",
+      },
+    ],
+  ],
+  schema: z.object({
+    amount: z
+      .number()
+      .positive()
+      .describe(
+        "The amount in tokens you'd like to withdraw from your drift account",
+      ),
+    symbol: z
+      .string()
+      .toUpperCase()
+      .describe("The symbol of the token you'd like to withdraw"),
+  }),
+  handler: async (agent, input) => {
+    try {
+      const tx = await withdrawFromDriftUserAccount(
+        agent,
+        input.amount,
+        input.symbol,
+      );
+
+      return {
+        status: "success",
+        message: "Funds withdrawn successfully",
+        signature: tx,
+      };
+    } catch (e) {
+      return {
+        status: "error",
+        // @ts-expect-error - error message is a string
+        message: `Failed to withdraw funds: ${e.message}`,
+      };
+    }
+  },
+};
+
+export default withdrawFromDriftAccountAction;

src/actions/index.ts

@@ -37,6 +37,11 @@ import requestWithdrawalFromVaultAction from "./drift/requestWithdrawalFromVault
 import withdrawFromVaultAction from "./drift/withdrawFromVault";
 import tradeDelegatedDriftVaultAction from "./drift/tradeDelegatedDriftVault";
 import vaultInfoAction from "./drift/vaultInfo";
+import createDriftUserAccountAction from "./drift/createDriftUserAccount";
+import tradeDriftPerpAccountAction from "./drift/tradePerpAccount";
+import doesUserHaveDriftAccountAction from "./drift/doesUserHaveDriftAccount";
+import depositToDriftUserAccountAction from "./drift/depositToDriftUserAccount";
+import withdrawFromDriftAccountAction from "./drift/withdrawFromDriftAccount";
 
 export const ACTIONS = {
   WALLET_ADDRESS_ACTION: getWalletAddressAction,
@@ -79,6 +84,11 @@ export const ACTIONS = {
   WITHDRAW_FROM_DRIFT_VAULT_ACTION: withdrawFromVaultAction,
   TRADE_DELEGATED_DRIFT_VAULT_ACTION: tradeDelegatedDriftVaultAction,
   DRIFT_VAULT_INFO_ACTION: vaultInfoAction,
+  CREATE_DRIFT_USER_ACCOUNT_ACTION: createDriftUserAccountAction,
+  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,
 };
 
 export type { Action, ActionExample, Handler } from "../types/action";

src/tools/drift.ts

@@ -0,0 +1,344 @@
+import {
+  BASE_PRECISION,
+  convertToNumber,
+  DRIFT_PROGRAM_ID,
+  DriftClient,
+  FastSingleTxSender,
+  getLimitOrderParams,
+  getMarketOrderParams,
+  getUserAccountPublicKeySync,
+  MainnetSpotMarkets,
+  numberToSafeBN,
+  PositionDirection,
+  PostOnlyParams,
+  PRICE_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";
+
+export async function initClients(agent: SolanaAgentKit) {
+  const wallet: IWallet = {
+    publicKey: agent.wallet.publicKey,
+    payer: agent.wallet,
+    signAllTransactions: async (txs) => {
+      for (const tx of txs) {
+        tx.sign(agent.wallet);
+      }
+      return txs;
+    },
+    signTransaction: async (tx) => {
+      tx.sign(agent.wallet);
+      return tx;
+    },
+  };
+
+  const driftClient = new DriftClient({
+    connection: agent.connection,
+    wallet,
+    env: "mainnet-beta",
+    txSender: new FastSingleTxSender({
+      connection: agent.connection,
+      wallet,
+      timeout: 30000,
+      blockhashRefreshInterval: 1000,
+      opts: {
+        commitment: agent.connection.commitment ?? "confirmed",
+        skipPreflight: false,
+        preflightCommitment: agent.connection.commitment ?? "confirmed",
+      },
+    }),
+  });
+  const vaultProgram = new anchor.Program(
+    IDL,
+    VAULT_PROGRAM_ID,
+    driftClient.provider,
+  );
+  const vaultClient = new VaultClient({
+    driftClient,
+    // @ts-expect-error - type mismatch due to different dep versions
+    program: vaultProgram,
+    cliMode: false,
+  });
+  await driftClient.subscribe();
+
+  async function cleanUp() {
+    await driftClient.unsubscribe();
+  }
+
+  return { driftClient, vaultClient, cleanUp };
+}
+
+/**
+ * Create a drift user account provided an amount
+ * @param amount amount of the token to deposit
+ * @param symbol symbol of the token to deposit
+ */
+export async function createDriftUserAccount(
+  agent: SolanaAgentKit,
+  amount: number,
+  symbol: string,
+) {
+  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();
+    const token = MainnetSpotMarkets.find(
+      (v) => v.symbol === symbol.toUpperCase(),
+    );
+
+    if (!token) {
+      throw new Error(`Token with symbol ${symbol} not found`);
+    }
+
+    if (!userAccountExists) {
+      const depositAmount = new BN(amount).mul(token.precision);
+      const [txSignature, account] =
+        await driftClient.initializeUserAccountAndDepositCollateral(
+          depositAmount,
+          getAssociatedTokenAddressSync(token.mint, agent.wallet.publicKey),
+        );
+
+      await cleanUp();
+      return { txSignature, account };
+    }
+
+    await cleanUp();
+    return {
+      message: "User account already exists",
+      account: user.userAccountPublicKey,
+    };
+  } catch (e) {
+    // @ts-expect-error - error message is a string
+    throw new Error(`Failed to create user account: ${e.message}`);
+  }
+}
+
+/**
+ * Deposit to your drift user account
+ * @param agent
+ * @param amount
+ * @param symbol
+ * @param address
+ * @returns
+ */
+export async function depositToDriftUserAccount(
+  agent: SolanaAgentKit,
+  amount: number,
+  symbol: string,
+  address?: string,
+) {
+  try {
+    const { driftClient, cleanUp } = await initClients(agent);
+    const publicKey = address ? new PublicKey(address) : agent.wallet.publicKey;
+    const user = new User({
+      driftClient,
+      userAccountPublicKey: getUserAccountPublicKeySync(
+        new PublicKey(DRIFT_PROGRAM_ID),
+        publicKey,
+      ),
+    });
+    const userAccountExists = await user.exists();
+    const token = MainnetSpotMarkets.find(
+      (v) => v.symbol === symbol.toUpperCase(),
+    );
+
+    if (!token) {
+      throw new Error(`Token with symbol ${symbol} not found`);
+    }
+
+    if (!userAccountExists) {
+      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),
+    );
+
+    await cleanUp();
+    return txSignature;
+  } catch (e) {
+    // @ts-expect-error - error message is a string
+    throw new Error(`Failed to deposit to user account: ${e.message}`);
+  }
+}
+
+export async function withdrawFromDriftUserAccount(
+  agent: SolanaAgentKit,
+  amount: number,
+  symbol: string,
+) {
+  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("You need to create a Drift user account first.");
+    }
+
+    const token = MainnetSpotMarkets.find(
+      (v) => v.symbol === symbol.toUpperCase(),
+    );
+
+    if (!token) {
+      throw new Error(`Token with symbol ${symbol} not found`);
+    }
+
+    const withdrawAmount = numberToSafeBN(amount, token.precision);
+
+    const txSignature = await driftClient.withdraw(
+      withdrawAmount,
+      token.marketIndex,
+      getAssociatedTokenAddressSync(token.mint, agent.wallet.publicKey),
+    );
+
+    await cleanUp();
+    return txSignature;
+  } catch (e) {
+    // @ts-expect-error - error message is a string
+    throw new Error(`Failed to withdraw from user account: ${e.message}`);
+  }
+}
+
+/**
+ * Open a perpetual trade on drift
+ * @param agent
+ * @param params.amount
+ * @param params.symbol
+ * @param params.action
+ * @param params.type
+ * @param params.price this should only be supplied if type is limit
+ * @param params.reduceOnly
+ */
+export async function driftPerpTrade(
+  agent: SolanaAgentKit,
+  params: {
+    amount: number;
+    symbol: string;
+    action: "long" | "short";
+    type: "market" | "limit";
+    price?: number;
+  },
+) {
+  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("You need to create a Drift user account first.");
+    }
+
+    const market = driftClient.getMarketIndexAndType(
+      `${params.symbol.toUpperCase()}-PERP`,
+    );
+
+    if (!market) {
+      throw new Error(`Token with symbol ${params.symbol} not found`);
+    }
+
+    const baseAssetPrice = driftClient.getOracleDataForPerpMarket(
+      market.marketIndex,
+    );
+    const convertedAmount =
+      params.amount / convertToNumber(baseAssetPrice.price, PRICE_PRECISION);
+
+    let signature: anchor.web3.TransactionSignature;
+
+    if (params.type === "limit") {
+      if (!params.price) {
+        throw new Error("Price is required for limit orders");
+      }
+
+      signature = await driftClient.placePerpOrder(
+        getLimitOrderParams({
+          baseAssetAmount: numberToSafeBN(convertedAmount, BASE_PRECISION),
+          reduceOnly: false,
+          direction:
+            params.action === "long"
+              ? PositionDirection.LONG
+              : PositionDirection.SHORT,
+          marketIndex: market.marketIndex,
+          price: numberToSafeBN(params.price, PRICE_PRECISION),
+          postOnly: PostOnlyParams.SLIDE,
+        }),
+      );
+    } else {
+      signature = await driftClient.placePerpOrder(
+        getMarketOrderParams({
+          baseAssetAmount: numberToSafeBN(convertedAmount, BASE_PRECISION),
+          reduceOnly: false,
+          direction:
+            params.action === "long"
+              ? PositionDirection.LONG
+              : PositionDirection.SHORT,
+          marketIndex: market.marketIndex,
+        }),
+      );
+    }
+
+    if (!signature) {
+      throw new Error("Failed to place order. Please make sure ");
+    }
+
+    await cleanUp();
+    return signature;
+  } catch (e) {
+    // @ts-expect-error - error message is a string
+    throw new Error(`Failed to place order: ${e.message}`);
+  }
+}
+
+/**
+ * Check if a user has a drift account
+ * @param agent
+ */
+export async function doesUserHaveDriftAccount(agent: SolanaAgentKit) {
+  try {
+    const { driftClient, cleanUp } = await initClients(agent);
+    const user = new User({
+      driftClient,
+      userAccountPublicKey: getUserAccountPublicKeySync(
+        new PublicKey(DRIFT_PROGRAM_ID),
+        agent.wallet.publicKey,
+      ),
+    });
+    user.getActivePerpPositions();
+    const userAccountExists = await user.exists();
+    await cleanUp();
+    return {
+      hasAccount: userAccountExists,
+      account: user.userAccountPublicKey,
+    };
+  } 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

@@ -1,8 +1,6 @@
 import {
   BASE_PRECISION,
   convertToNumber,
-  DriftClient,
-  FastSingleTxSender,
   getLimitOrderParams,
   getMarketOrderParams,
   getOrderParams,
@@ -16,17 +14,12 @@ import {
   PRICE_PRECISION,
   QUOTE_PRECISION,
   TEN,
-  type IWallet,
 } from "@drift-labs/sdk";
 import {
-  VAULT_PROGRAM_ID,
-  VaultClient,
-  IDL,
   WithdrawUnit,
   encodeName,
   getVaultDepositorAddressSync,
 } from "@drift-labs/vaults-sdk";
-import * as anchor from "@coral-xyz/anchor";
 import {
   ComputeBudgetProgram,
   PublicKey,
@@ -34,6 +27,7 @@ import {
 } from "@solana/web3.js";
 import type { SolanaAgentKit } from "../agent";
 import { BN } from "bn.js";
+import { initClients } from "./drift";
 
 export function getMarketIndexAndType(name: `${string}-${string}`) {
   const [symbol, type] = name.toUpperCase().split("-");
@@ -53,58 +47,6 @@ export function getMarketIndexAndType(name: `${string}-${string}`) {
   return { marketIndex: token.marketIndex, marketType: MarketType.SPOT };
 }
 
-async function initClients(agent: SolanaAgentKit) {
-  const wallet: IWallet = {
-    publicKey: agent.wallet.publicKey,
-    payer: agent.wallet,
-    signAllTransactions: async (txs) => {
-      for (const tx of txs) {
-        tx.sign(agent.wallet);
-      }
-      return txs;
-    },
-    signTransaction: async (tx) => {
-      tx.sign(agent.wallet);
-      return tx;
-    },
-  };
-
-  const driftClient = new DriftClient({
-    connection: agent.connection,
-    wallet,
-    env: "mainnet-beta",
-    txSender: new FastSingleTxSender({
-      connection: agent.connection,
-      wallet,
-      timeout: 30000,
-      blockhashRefreshInterval: 1000,
-      opts: {
-        commitment: agent.connection.commitment ?? "confirmed",
-        skipPreflight: false,
-        preflightCommitment: agent.connection.commitment ?? "confirmed",
-      },
-    }),
-  });
-  const vaultProgram = new anchor.Program(
-    IDL,
-    VAULT_PROGRAM_ID,
-    driftClient.provider,
-  );
-  const vaultClient = new VaultClient({
-    driftClient,
-    // @ts-expect-error - type mismatch due to different dep versions
-    program: vaultProgram,
-    cliMode: false,
-  });
-  await driftClient.subscribe();
-
-  async function cleanUp() {
-    await driftClient.unsubscribe();
-  }
-
-  return { driftClient, vaultClient, cleanUp };
-}
-
 async function getOrCreateVaultDepositor(agent: SolanaAgentKit, vault: string) {
   const { vaultClient, cleanUp } = await initClients(agent);
   const vaultPublicKey = new PublicKey(vault);
@@ -236,27 +178,49 @@ export async function updateVault(
   },
 ) {
   try {
-    const { vaultClient, cleanUp } = await initClients(agent);
+    const { vaultClient, cleanUp, driftClient } = await initClients(agent);
     const vaultPublicKey = new PublicKey(vault);
     const vaultDetails = await vaultClient.getVault(vaultPublicKey);
 
+    const spotMarket = driftClient.getSpotMarketAccount(
+      vaultDetails.spotMarketIndex,
+    );
+
+    if (!spotMarket) {
+      throw new Error("Market not found");
+    }
+
+    const spotPrecision = TEN.pow(new BN(spotMarket.decimals));
+
     const tx = await vaultClient.managerUpdateVault(vaultPublicKey, {
       redeemPeriod: new BN(
         params.redeemPeriod
           ? params.redeemPeriod * 86400
           : vaultDetails.redeemPeriod,
       ),
-      maxTokens: new BN(params.maxTokens ?? vaultDetails.maxTokens),
-      minDepositAmount: new BN(
-        params.minDepositAmount ?? vaultDetails.minDepositAmount,
-      ),
-      managementFee: new BN(params.managementFee ?? vaultDetails.managementFee),
-      profitShare: new BN(
-        params.profitShare ?? vaultDetails.profitShare,
-      ).toNumber(),
-      hurdleRate: new BN(
-        params.hurdleRate ?? vaultDetails.hurdleRate,
-      ).toNumber(),
+      maxTokens: params.maxTokens
+        ? numberToSafeBN(params.maxTokens, spotPrecision)
+        : vaultDetails.maxTokens,
+      minDepositAmount: params.minDepositAmount
+        ? numberToSafeBN(params.minDepositAmount, spotPrecision)
+        : vaultDetails.minDepositAmount,
+      managementFee: params.managementFee
+        ? new BN(params.managementFee)
+            .mul(PERCENTAGE_PRECISION)
+            .div(new BN(100))
+        : vaultDetails.managementFee,
+      profitShare: params.profitShare
+        ? new BN(params.profitShare)
+            .mul(PERCENTAGE_PRECISION)
+            .div(new BN(100))
+            .toNumber()
+        : vaultDetails.profitShare,
+      hurdleRate: params.hurdleRate
+        ? new BN(params.hurdleRate)
+            .mul(PERCENTAGE_PRECISION)
+            .div(new BN(100))
+            .toNumber()
+        : vaultDetails.hurdleRate,
       permissioned: params.permissioned ?? vaultDetails.permissioned,
     });
 
@@ -494,15 +458,18 @@ export async function tradeDriftVault(
     const perpMarketIndexAndType = getMarketIndexAndType(
       `${symbol.toUpperCase()}-PERP`,
     );
+    const perpMarketAccount = driftClient.getPerpMarketAccount(
+      perpMarketIndexAndType.marketIndex,
+    );
 
-    if (!perpMarketIndexAndType) {
+    if (!perpMarketIndexAndType || !perpMarketAccount) {
       throw new Error(
         "Invalid symbol: Drift doesn't have a market for this token",
       );
     }
 
     const perpOracle = driftClient.getOracleDataForPerpMarket(
-      perpMarketIndexAndType.marketIndex,
+      perpMarketAccount.marketIndex,
     );
     const oraclePriceNumber = convertToNumber(
       perpOracle.price,
@@ -530,7 +497,7 @@ export async function tradeDriftVault(
               action === "buy"
                 ? PositionDirection.LONG
                 : PositionDirection.SHORT,
-            marketIndex: perpMarketIndexAndType.marketIndex,
+            marketIndex: perpMarketAccount.marketIndex,
             postOnly: PostOnlyParams.SLIDE,
           }),
         ),
@@ -548,7 +515,7 @@ export async function tradeDriftVault(
               action === "buy"
                 ? PositionDirection.LONG
                 : PositionDirection.SHORT,
-            marketIndex: perpMarketIndexAndType.marketIndex,
+            marketIndex: perpMarketAccount.marketIndex,
           }),
         ),
       ]);

src/tools/index.ts

@@ -52,3 +52,4 @@ export * from "./flash_close_trade";
 export * from "./create_3land_collectible";
 
 export * from "./drift_vault";
+export * from "./drift";