fix: drift user account balances and position amounts (#232)

# Pull Request Description

This PR fixes a drift account info issue pointed out by the Drift
protocol team

## Changes Made
This PR adds the following changes:
<!-- List the key changes made in this PR -->
- I updated the drift and drift vault SDK versions
- I corrected the `driftUserAccountInfo` function

## Prompt executed by agent 
<!-- If applicable, provide example usage, transactions, or screenshots
-->
Example transaction: 
<img width="592" alt="Screenshot 2025-01-23 at 09 11 58"
src="https://github.com/user-attachments/assets/24c48204-de50-4231-bc73-abfe04776d51"
/>


## Prompt Used
<!-- If relevant, include the prompt or configuration used -->
```
show my drift account info
```

## Additional Notes
<!-- Any additional information that reviewers should know -->

## Checklist
- [x] I have tested these changes locally
- [ ] I have updated the documentation
- [ ] I have added a transaction link
- [x] I have added the prompt used to test it

Author: arihantbansal

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",

pnpm-lock.yaml

@@ -1,6 +1,7 @@
 import {
   BASE_PRECISION,
   BigNum,
+  BulkAccountLoader,
   calculateDepositRate,
   calculateEstimatedEntryPriceWithL2,
   calculateInterestRate,
@@ -14,7 +15,9 @@ import {
   getInsuranceFundStakeAccountPublicKey,
   getLimitOrderParams,
   getMarketOrderParams,
+  getTokenAmount,
   getUserAccountPublicKeySync,
+  isVariant,
   JupiterClient,
   MainnetPerpMarkets,
   MainnetSpotMarkets,
@@ -71,6 +74,10 @@ export async function initClients(
     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 +446,14 @@ export async function doesUserHaveDriftAccount(agent: SolanaAgentKit) {
 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 +462,80 @@ export async function driftUserAccountInfo(agent: SolanaAgentKit) {
     }
     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