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

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}`);
+  }
+}