mirror of
https://github.com/d0zingcat/solana-agent-kit.git
synced 2026-05-13 23:16:55 +00:00
447 lines
12 KiB
TypeScript
447 lines
12 KiB
TypeScript
import {
|
|
BASE_PRECISION,
|
|
convertToNumber,
|
|
DRIFT_PROGRAM_ID,
|
|
DriftClient,
|
|
FastSingleTxSender,
|
|
getLimitOrderParams,
|
|
getMarketOrderParams,
|
|
getUserAccountPublicKeySync,
|
|
MainnetSpotMarkets,
|
|
numberToSafeBN,
|
|
PositionDirection,
|
|
PostOnlyParams,
|
|
PRICE_PRECISION,
|
|
QUOTE_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 { getAssociatedTokenAddressSync } from "@solana/spl-token";
|
|
import { PublicKey } from "@solana/web3.js";
|
|
import { Transaction } from "@solana/web3.js";
|
|
import { ComputeBudgetProgram } 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 = numberToSafeBN(amount, 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 isRepay
|
|
* @returns
|
|
*/
|
|
export async function depositToDriftUserAccount(
|
|
agent: SolanaAgentKit,
|
|
amount: number,
|
|
symbol: string,
|
|
isRepay = false,
|
|
) {
|
|
try {
|
|
const { driftClient, cleanUp } = await initClients(agent);
|
|
const publicKey = 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 = numberToSafeBN(amount, token.precision);
|
|
|
|
const [depInstruction, latestBlockhash] = await Promise.all([
|
|
driftClient.getDepositTxnIx(
|
|
depositAmount,
|
|
token.marketIndex,
|
|
getAssociatedTokenAddressSync(token.mint, publicKey),
|
|
undefined,
|
|
isRepay,
|
|
),
|
|
driftClient.connection.getLatestBlockhash(),
|
|
]);
|
|
|
|
const tx = new Transaction().add(...depInstruction).add(
|
|
ComputeBudgetProgram.setComputeUnitPrice({
|
|
microLamports: 0.000001 * 1000000 * 1000000,
|
|
}),
|
|
);
|
|
tx.recentBlockhash = latestBlockhash.blockhash;
|
|
tx.sign(agent.wallet);
|
|
const txSignature = await driftClient.txSender.sendRawTransaction(
|
|
tx.serialize(),
|
|
{ ...driftClient.opts },
|
|
);
|
|
|
|
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,
|
|
isBorrow = false,
|
|
) {
|
|
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 [withdrawInstruction, latestBlockhash] = await Promise.all([
|
|
driftClient.getWithdrawalIxs(
|
|
withdrawAmount,
|
|
token.marketIndex,
|
|
getAssociatedTokenAddressSync(token.mint, agent.wallet.publicKey),
|
|
!isBorrow,
|
|
),
|
|
driftClient.connection.getLatestBlockhash(),
|
|
]);
|
|
|
|
const tx = new Transaction().add(...withdrawInstruction).add(
|
|
ComputeBudgetProgram.setComputeUnitPrice({
|
|
microLamports: 0.000001 * 1000000 * 1000000,
|
|
}),
|
|
);
|
|
tx.recentBlockhash = latestBlockhash.blockhash;
|
|
tx.sign(agent.wallet);
|
|
|
|
const txSignature = await driftClient.txSender.sendRawTransaction(
|
|
tx.serialize(),
|
|
{ ...driftClient.opts },
|
|
);
|
|
|
|
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,
|
|
}),
|
|
{
|
|
computeUnitsPrice: 0.000001 * 1000000 * 1000000,
|
|
},
|
|
);
|
|
} else {
|
|
signature = await driftClient.placePerpOrder(
|
|
getMarketOrderParams({
|
|
baseAssetAmount: numberToSafeBN(convertedAmount, BASE_PRECISION),
|
|
reduceOnly: false,
|
|
direction:
|
|
params.action === "long"
|
|
? PositionDirection.LONG
|
|
: PositionDirection.SHORT,
|
|
marketIndex: market.marketIndex,
|
|
}),
|
|
{
|
|
computeUnitsPrice: 0.000001 * 1000000 * 1000000,
|
|
},
|
|
);
|
|
}
|
|
|
|
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}`);
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Get account info for a drift User
|
|
* @param agent
|
|
* @returns
|
|
*/
|
|
export async function driftUserAccountInfo(agent: SolanaAgentKit) {
|
|
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("User account does not exist");
|
|
}
|
|
await user.subscribe();
|
|
const account = user.getUserAccount();
|
|
await user.unsubscribe();
|
|
|
|
await cleanUp();
|
|
const perpPositions = account.perpPositions.map((pos) => ({
|
|
...pos,
|
|
baseAssetAmount: convertToNumber(pos.baseAssetAmount, BASE_PRECISION),
|
|
settledPnl: convertToNumber(pos.settledPnl, QUOTE_PRECISION),
|
|
}));
|
|
const spotPositions = account.spotPositions.map((pos) => ({
|
|
...pos,
|
|
scaledBalance: convertToNumber(pos.scaledBalance, BASE_PRECISION),
|
|
cumulativeDeposits: convertToNumber(
|
|
pos.cumulativeDeposits,
|
|
BASE_PRECISION,
|
|
),
|
|
symbol: MainnetSpotMarkets.find((v) => v.marketIndex === pos.marketIndex)
|
|
?.symbol,
|
|
}));
|
|
|
|
return {
|
|
...account,
|
|
name: account.name,
|
|
authority: account.authority,
|
|
totalDeposits: `$${convertToNumber(account.totalDeposits, QUOTE_PRECISION)}`,
|
|
totalWithdraws: `$${convertToNumber(account.totalWithdraws, QUOTE_PRECISION)}`,
|
|
settledPerpPnl: `$${convertToNumber(account.settledPerpPnl, QUOTE_PRECISION)}`,
|
|
lastActiveSlot: account.lastActiveSlot.toNumber(),
|
|
perpPositions,
|
|
spotPositions,
|
|
};
|
|
} catch (e) {
|
|
// @ts-expect-error - error message is a string
|
|
throw new Error(`Failed to check user account: ${e.message}`);
|
|
}
|
|
}
|