feat + fix: add action to get user drift account info and fix deposit, borrow, lend and withdraw on drift

src/actions/drift/depositToDriftUserAccount.ts

@@ -10,6 +10,7 @@ const depositToDriftUserAccountAction: Action = {
     "deposit into drift user account",
     "add funds to drift user account",
     "add funds to my drift account",
+    "deposit collateral into drift account",
   ],
   examples: [
     [
@@ -27,22 +28,6 @@ const depositToDriftUserAccountAction: Action = {
         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
@@ -55,7 +40,11 @@ const depositToDriftUserAccountAction: Action = {
       .string()
       .toUpperCase()
       .describe("The symbol of the token you'd like to deposit"),
-    address: z.string().optional().describe("The drift user account address"),
+    repay: z
+      .boolean()
+      .optional()
+      .default(false)
+      .describe("Whether or not to repay the borrowed funds in the account"),
   }),
   handler: async (agent: SolanaAgentKit, input) => {
     try {
@@ -63,7 +52,7 @@ const depositToDriftUserAccountAction: Action = {
         agent,
         input.amount as number,
         input.symbol as string,
-        input.address,
+        input.repay as boolean,
       );
 
       return {

src/actions/drift/driftUserAccountInfo.ts

@@ -0,0 +1,39 @@
+import { z } from "zod";
+import type { Action } from "../../types";
+import { driftUserAccountInfo } from "../../tools";
+
+const driftUserAccountInfoAction: Action = {
+  name: "DRIFT_USER_ACCOUNT_INFO",
+  similes: ["get drift user account info", "get drift account info"],
+  description: "Get information about your drift account",
+  examples: [
+    [
+      {
+        input: {},
+        explanation: "Get information about your drift account",
+        output: {
+          status: "success",
+          data: {},
+        },
+      },
+    ],
+  ],
+  schema: z.object({}),
+  handler: async (agent) => {
+    try {
+      const accountInfo = await driftUserAccountInfo(agent);
+      return {
+        status: "success",
+        data: accountInfo,
+      };
+    } catch (e) {
+      return {
+        status: "error",
+        // @ts-expect-error - error message is a string
+        message: `Failed to get drift account info: ${e.message}`,
+      };
+    }
+  },
+};
+
+export default driftUserAccountInfoAction;

src/actions/drift/tradePerpAccount.ts

@@ -47,7 +47,10 @@ export const tradeDriftPerpAccountAction: Action = {
   ],
   schema: z.object({
     amount: z.number().positive(),
-    symbol: z.string().min(3).max(10),
+    symbol: z
+      .string()
+      .toUpperCase()
+      .describe("Symbol of the token to open a position on "),
     action: z.enum(["long", "short"]),
     type: z.enum(["market", "limit"]),
     price: z.number().positive().optional(),

src/actions/drift/withdrawFromDriftAccount.ts

@@ -3,12 +3,16 @@ import type { Action } from "../../types";
 import { withdrawFromDriftUserAccount } from "../../tools";
 
 const withdrawFromDriftAccountAction: Action = {
-  name: "WITHDRAW_FROM_DRIFT_ACCOUNT",
+  name: "WITHDRAW_OR_BORROW_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",
+    "borrow from drift account",
+    "borrow funds from my drift account",
+    "borrow from drift",
+    "withdraw from drift",
   ],
   examples: [
     [
@@ -38,6 +42,13 @@ const withdrawFromDriftAccountAction: Action = {
       .string()
       .toUpperCase()
       .describe("The symbol of the token you'd like to withdraw"),
+    isBorrow: z
+      .boolean()
+      .optional()
+      .default(false)
+      .describe(
+        "Whether or not to borrow funds based on collateral provided instead of withdrawing",
+      ),
   }),
   handler: async (agent, input) => {
     try {
@@ -45,6 +56,7 @@ const withdrawFromDriftAccountAction: Action = {
         agent,
         input.amount,
         input.symbol,
+        input.isBorrow,
       );
 
       return {

src/actions/index.ts

@@ -42,6 +42,7 @@ import tradeDriftPerpAccountAction from "./drift/tradePerpAccount";
 import doesUserHaveDriftAccountAction from "./drift/doesUserHaveDriftAccount";
 import depositToDriftUserAccountAction from "./drift/depositToDriftUserAccount";
 import withdrawFromDriftAccountAction from "./drift/withdrawFromDriftAccount";
+import driftUserAccountInfoAction from "./drift/driftUserAccountInfo";
 
 export const ACTIONS = {
   WALLET_ADDRESS_ACTION: getWalletAddressAction,
@@ -88,7 +89,8 @@ export const ACTIONS = {
   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,
+  WITHDRAW_OR_BORROW_FROM_DRIFT_ACCOUNT_ACTION: withdrawFromDriftAccountAction,
+  DRIFT_USER_ACCOUNT_INFO_ACTION: driftUserAccountInfoAction,
 };
 
 export type { Action, ActionExample, Handler } from "../types/action";