Merge branch 'main' into feat/mayan-swap-sdk

* main: chore: remove unused import Update README.md: change 15+ actions to 60+ actions do not log kp fix: spot position balance signs fix: spot position balance type and signs fix: overall drift account usd balance fix: spot position balance type fix: perp position precision fix: drift user account balances and position amounts
sendaifun · Jan 28, 2025 · e846b10 · e846b10
2 parents 7e5afdc + 0e8c1d1
commit e846b10
Show file tree

Hide file tree

Showing 5 changed files with 92 additions and 39 deletions.
diff --git a/README.md b/README.md
@@ -10,7 +10,7 @@
 
 </div>
 
-An open-source toolkit for connecting AI agents to Solana protocols. Now, any agent, using any model can autonomously perform 15+ Solana actions:
+An open-source toolkit for connecting AI agents to Solana protocols. Now, any agent, using any model can autonomously perform 60+ Solana actions:
 
 - Trade tokens
 - Launch new tokens

diff --git a/package.json b/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",

diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml
diff --git a/src/tools/drift/drift.ts b/src/tools/drift/drift.ts
@@ -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

diff --git a/src/utils/keypair.ts b/src/utils/keypair.ts
@@ -4,13 +4,9 @@ import {
   Transaction,
   VersionedTransaction,
 } from "@solana/web3.js";
-import bs58 from "bs58";
 
 export const keypair = Keypair.generate();
 
-console.log(keypair.publicKey.toString());
-console.log(bs58.encode(keypair.secretKey));
-
 export class Wallet {
   private _signer: Keypair;