fix: drift user account balances and position amounts

package.json

@@ -35,8 +35,8 @@
     "@bonfida/spl-name-service": "^3.0.7",
     "@cks-systems/manifest-sdk": "0.1.59",
     "@coral-xyz/anchor": "0.29",
-    "@drift-labs/sdk": "2.107.0-beta.3",
-    "@drift-labs/vaults-sdk": "^0.2.49",
+    "@drift-labs/sdk": "2.108.0-beta.4",
+    "@drift-labs/vaults-sdk": "^0.3.2",
     "@langchain/core": "^0.3.26",
     "@langchain/groq": "^0.1.2",
     "@langchain/langgraph": "^0.2.36",

src/tools/drift/drift.ts

@@ -1,6 +1,7 @@
 import {
   BASE_PRECISION,
   BigNum,
+  BulkAccountLoader,
   calculateDepositRate,
   calculateEstimatedEntryPriceWithL2,
   calculateInterestRate,
@@ -14,7 +15,9 @@
   getInsuranceFundStakeAccountPublicKey,
   getLimitOrderParams,
   getMarketOrderParams,
+  getTokenAmount,
   getUserAccountPublicKeySync,
+  isVariant,
   JupiterClient,
   MainnetPerpMarkets,
   MainnetSpotMarkets,
@@ -24,6 +27,7 @@
   PostOnlyParams,
   PRICE_PRECISION,
   QUOTE_PRECISION,
+  SpotBalanceType,
   User,
   type IWallet,
 } from "@drift-labs/sdk";
@@ -71,6 +75,10 @@
     txParams: {
       computeUnitsPrice: MINIMUM_COMPUTE_PRICE_FOR_COMPLEX_ACTIONS,
     },
+    accountSubscription: {
+      type: "polling",
+      accountLoader: new BulkAccountLoader(agent.connection, "processed", 10),
+    },
     txSender: new FastSingleTxSender({
       connection: agent.connection,
       wallet,
@@ -439,12 +447,14 @@
 export async function driftUserAccountInfo(agent: SolanaAgentKit) {
   try {
     const { driftClient, cleanUp } = await initClients(agent);
+    const userPublicKey = getUserAccountPublicKeySync(
+      new PublicKey(DRIFT_PROGRAM_ID),
+      agent.wallet.publicKey,
+    );
+
     const user = new User({
       driftClient,
-      userAccountPublicKey: getUserAccountPublicKeySync(
-        new PublicKey(DRIFT_PROGRAM_ID),
-        agent.wallet.publicKey,
-      ),
+      userAccountPublicKey: userPublicKey,
     });
     const userAccountExists = await user.exists();
 
@@ -453,32 +463,80 @@
     }
     await user.subscribe();
     const account = user.getUserAccount();
-    await user.unsubscribe();
 
-    await cleanUp();
     const perpPositions = account.perpPositions.map((pos) => ({
-      ...pos,
+      market: MainnetPerpMarkets[pos.marketIndex].symbol,
       baseAssetAmount: convertToNumber(pos.baseAssetAmount, BASE_PRECISION),
+      quoteAssetAmount: convertToNumber(
+        pos.quoteAssetAmount.abs(),
+        QUOTE_PRECISION,
+      ),
+      quoteEntryAmount: convertToNumber(
+        pos.quoteEntryAmount.abs(),
+        QUOTE_PRECISION,
+      ),
+      quoteBreakEvenAmount: convertToNumber(
+        pos.quoteBreakEvenAmount.abs(),
+        QUOTE_PRECISION,
+      ),
       settledPnl: convertToNumber(pos.settledPnl, QUOTE_PRECISION),
+      openAsks: pos.openAsks.toNumber(),
+      openBids: pos.openBids.toNumber(),
+      openOrders: pos.openOrders,
+      positionType:
+        convertToNumber(pos.baseAssetAmount, BASE_PRECISION) > 0
+          ? "long"
+          : "short",
     }));
-    const spotPositions = account.spotPositions.map((pos) => ({
-      ...pos,
-      availableBalance: convertToNumber(
+    const spotPositions = account.spotPositions.map((pos) => {
+      const spotMarketAccount = driftClient.getSpotMarketAccount(
+        pos.marketIndex,
+      );
+
+      if (!spotMarketAccount) {
+        return;
+      }
+
+      const tokenBalance = getTokenAmount(
         pos.scaledBalance,
-        MainnetSpotMarkets[pos.marketIndex].precision,
-      ),
-      symbol: MainnetSpotMarkets.find((v) => v.marketIndex === pos.marketIndex)
-        ?.symbol,
-    }));
+        spotMarketAccount,
+        pos.balanceType,
+      );
+
+      return {
+        availableBalance:
+          (isVariant(pos.balanceType, "borrow") ? -1 : 1) *
+          convertToNumber(
+            tokenBalance,
+            MainnetSpotMarkets[pos.marketIndex].precision,
+          ),
+        symbol: MainnetSpotMarkets[pos.marketIndex].symbol,
+        openAsks: pos.openAsks.toNumber(),
+        openBids: pos.openBids.toNumber(),
+        openOrders: pos.openOrders,
+        type: isVariant(pos.balanceType, "borrow") ? "borrow" : "deposit",
+      };
+    });
+
+    const overallUserBalance = user.getNetSpotMarketValue();
+    const unrealizedPnl = user.getUnrealizedPNL(true);
+    const netUSDValue = convertToNumber(
+      overallUserBalance.add(unrealizedPnl),
+      QUOTE_PRECISION,
+    );
+
+    await cleanUp();
+    await user.unsubscribe();
 
     return {
-      ...account,
       name: account.name,
+      accountAddress: userPublicKey.toBase58(),
       authority: account.authority,
+      overallBalance: netUSDValue,
       settledPerpPnl: `$${convertToNumber(account.settledPerpPnl, QUOTE_PRECISION)}`,
       lastActiveSlot: account.lastActiveSlot.toNumber(),
-      perpPositions,
-      spotPositions,
+      perpPositions: perpPositions.filter((pos) => pos.baseAssetAmount !== 0),
+      spotPositions: spotPositions.filter((pos) => pos?.availableBalance !== 0),
     };
   } catch (e) {
     // @ts-expect-error - error message is a string