@shogun-sdk/money-legos/src/lib/jito.ts

Shogun Money Legos: clients and types for quotes, memes, prices, balances, fees, validations, etc.

import {
  AddressLookupTableAccount,
  ComputeBudgetProgram,
  SystemProgram,
  TransactionMessage,
  VersionedTransaction,
} from '@solana/web3.js';
import { SOLANA_CHAIN_ID } from '../config/chains.js';
import type { QuoteTypes } from '../types/index.js';
import { getSolanaProvider } from '../utils/rpc.js';

export const PRIORITY_FEE_SHARE = 0.7;
export const DEFAULT_JITO_TIP = 6000000n; // in lamports: 0.006 SOL

export async function processSolana(quote: QuoteTypes, jitoTip = DEFAULT_JITO_TIP) {
  const solanaTransactions = quote.inputAmount.chainId === SOLANA_CHAIN_ID;
  if (Array.isArray(quote.calldatas) && solanaTransactions) {
    // https://docs.jito.wtf/lowlatencytxnsend/#tip-amount
    if (quote.calldatas.length === 1 && quote.calldatas[0]) {
      // Jito support guy said that for jito bundle there's no reason to use priority fees
      // adding/setting priority fee instructions to a single Solana transaction
      const transaction = VersionedTransaction.deserialize(Buffer.from(quote.calldatas[0].data, 'base64'));
      const { transaction: priorityFeeTransaction, jitoTipsTotal } = await decompileAndUpdateSingleJitoTipTransaction(
        transaction,
        Number(jitoTip),
      );
      quote.calldatas[0]['data'] = Buffer.from(priorityFeeTransaction.serialize()).toString('base64');
      quote.jitoTipsTotal = jitoTipsTotal;
    }
  }

  return quote;
}

export async function decompileAndUpdateSingleJitoTipTransaction(
  transaction: VersionedTransaction,
  jitoTip?: number | string,
): Promise<{ transaction: VersionedTransaction; jitoTipsTotal: bigint }> {
  const provider = await getSolanaProvider();

  // Decompile transaction to get address lookup table accounts
  const addressLookupTableAccounts = await Promise.all(
    transaction.message.addressTableLookups.map(async (lookup) => {
      const accountInfo = await provider.getAccountInfo(lookup.accountKey);
      if (!accountInfo) {
        throw new Error(`Failed to fetch account info for ${lookup.accountKey.toString()}`);
      }
      const state = AddressLookupTableAccount.deserialize(accountInfo.data);
      return new AddressLookupTableAccount({ key: lookup.accountKey, state });
    }),
  );

  let message = TransactionMessage.decompile(transaction.message, { addressLookupTableAccounts });
  const { message: updatedMessage, jitoTipsTotal } = updateInstructionsForSingleJitoTransaction(
    message,
    Number(jitoTip),
  );
  message = updatedMessage;

  // Compile the modified message back into a VersionedTransaction
  const compiledMessage = message.compileToV0Message(addressLookupTableAccounts);
  return {
    transaction: new VersionedTransaction(compiledMessage),
    jitoTipsTotal,
  };
}

export function updateInstructionsForSingleJitoTransaction(
  message: TransactionMessage,
  jitoTip: number,
): { message: TransactionMessage; jitoTipsTotal: bigint } {
  let jitoTipsTotal = BigInt(0);

  if (jitoTip === 0 || isNaN(jitoTip)) return { message, jitoTipsTotal };

  console.log('>>> Jito check: ', jitoTip);
  const computeUnitsMicroLamports = Math.ceil(PRIORITY_FEE_SHARE * jitoTip);
  const newJitoTip = Math.ceil(jitoTip - computeUnitsMicroLamports);
  // Create the priority fee instruction
  const priorityFeeIx = ComputeBudgetProgram.setComputeUnitPrice({
    microLamports: computeUnitsMicroLamports,
  });

  let computeUnitPriceUpdated = false;

  // Update or insert the `setComputeUnitPrice` instruction
  for (let i = 0; i < message.instructions.length; i++) {
    const instruction = message.instructions[i];

    if (
      instruction &&
      instruction.programId.equals(ComputeBudgetProgram.programId) &&
      instruction.data[0] === 3 // setComputeUnitPrice opcode
    ) {
      console.log('>>> Compute unit price updated to', computeUnitsMicroLamports);
      message.instructions[i] = priorityFeeIx;
      computeUnitPriceUpdated = true;
      break;
    }
  }

  // If no existing compute unit price instruction was found, add it at the beginning
  if (!computeUnitPriceUpdated) {
    console.log('>>> Adding compute unit price instruction:', computeUnitsMicroLamports);
    message.instructions.unshift(priorityFeeIx);
  }

  // Find and update the Jito tip instruction
  for (let i = message.instructions.length - 1; i >= 0; i--) {
    const instruction = message.instructions[i];

    if (instruction?.programId.equals(SystemProgram.programId) && instruction?.data.length === 8) {
      // Decode the existing lamports amount
      const oldTipAmount = instruction.data.readBigUInt64LE(0);

      if (oldTipAmount === BigInt(jitoTip) && instruction.keys?.[0] && instruction.keys?.[1]) {
        console.log('>>> Jito tip updated to', newJitoTip);

        jitoTipsTotal += BigInt(newJitoTip);
        // Replace the instruction with a new one containing the updated amount
        message.instructions[i] = SystemProgram.transfer({
          fromPubkey: instruction.keys[0].pubkey,
          toPubkey: instruction.keys[1].pubkey,
          lamports: newJitoTip,
        });

        break;
      }
    }
  }

  return { message, jitoTipsTotal };
}