feat: handle shapeshift fees for jupiter and rate balance check

package.json

@@ -73,7 +73,6 @@
     "@chakra-ui/react": "^2.8.2",
     "@chakra-ui/system": "^2.6.2",
     "@chakra-ui/tag": "^3.1.1",
-    "@coral-xyz/anchor": "0.29.0",
     "@cowprotocol/app-data": "^2.3.0",
     "@emotion/react": "^11.13.0",
     "@emotion/styled": "^11.13.0",

packages/swapper/package.json

@@ -17,6 +17,7 @@
     "dev": "yarn run -T tsc --build --watch"
   },
   "dependencies": {
+    "@coral-xyz/anchor": "0.29.0",
     "@shapeshiftoss/caip": "workspace:^",
     "@shapeshiftoss/chain-adapters": "workspace:^",
     "@shapeshiftoss/types": "workspace:^",

packages/swapper/src/swappers/JupiterSwapper/swapperApi/getTradeQuote.ts

@@ -33,6 +33,7 @@ import { referralIdl } from '../idls/referral'
 import {
   JUPITER_AFFILIATE_CONTRACT_ADDRESS,
   JUPITER_COMPUTE_UNIT_MARGIN_MULTIPLIER,
+  PDA_ACCOUNT_CREATION_COST,
   SHAPESHIFT_JUPITER_REFERRAL_KEY,
 } from '../utils/constants'
 import {
@@ -61,6 +62,8 @@ export const getTradeQuote = async (
 
   const jupiterUrl = deps.config.REACT_APP_JUPITER_API_URL
 
+  const solAsset = assetsById[solAssetId]
+
   if (accountNumber === undefined) {
     return Err(
       makeSwapErrorRight({
@@ -99,6 +102,15 @@ export const getTradeQuote = async (
     )
   }
 
+  if (!solAsset) {
+    return Err(
+      makeSwapErrorRight({
+        message: `solAsset is required`,
+        code: TradeQuoteError.UnknownError,
+      }),
+    )
+  }
+
   const maybePriceResponse = await getJupiterPrice({
     apiUrl: jupiterUrl,
     sourceAsset: sellAsset.assetId === solAssetId ? wrappedSolAssetId : sellAsset.assetId,
@@ -182,7 +194,8 @@ export const getTradeQuote = async (
       sellAssetAccount,
       programId: new PublicKey(JUPITER_AFFILIATE_CONTRACT_ADDRESS),
       instructionData,
-      mint: buyAssetAddress,
+      buyTokenId: buyAssetAddress,
+      sellTokenId: sellAssetAddress,
       connection: adapter.getConnection(),
     })
 
@@ -193,7 +206,7 @@ export const getTradeQuote = async (
     rawQuote: priceResponse,
     // Shared account is not supported for simple AMMs
     useSharedAccounts: priceResponse.routePlan.length > 1 && isCrossAccountTrade ? true : false,
-    feeAccount: tokenAccount.toString(),
+    feeAccount: affiliateBps !== '0' ? tokenAccount?.toString() : undefined,
   })
 
   if (maybeSwapResponse.isErr()) {
@@ -279,6 +292,16 @@ export const getTradeQuote = async (
     {} as Record<AssetId, ProtocolFee>,
   )
 
+  if (feeAccountInstruction) {
+    protocolFees[solAssetId] = {
+      requiresBalance: true,
+      amountCryptoBaseUnit: bnOrZero(protocolFees[solAssetId]?.amountCryptoBaseUnit)
+        .plus(PDA_ACCOUNT_CREATION_COST)
+        .toFixed(),
+      asset: solAsset,
+    }
+  }
+
   const quotes: TradeQuote[] = []
 
   const feeData = await getFeeData()

packages/swapper/src/swappers/JupiterSwapper/swapperApi/getTradeRate.ts

@@ -1,3 +1,4 @@
+import { BorshInstructionCoder } from '@coral-xyz/anchor'
 import type { AssetId } from '@shapeshiftoss/caip'
 import {
   ASSET_NAMESPACE,
@@ -13,6 +14,7 @@ import type { KnownChainIds } from '@shapeshiftoss/types'
 import { bn, bnOrZero, convertDecimalPercentageToBasisPoints } from '@shapeshiftoss/utils'
 import type { Result } from '@sniptt/monads'
 import { Err, Ok } from '@sniptt/monads'
+import { PublicKey } from '@solana/web3.js'
 import { v4 as uuid } from 'uuid'
 
 import type {
@@ -24,8 +26,18 @@ import type {
 } from '../../../types'
 import { SwapperName, TradeQuoteError } from '../../../types'
 import { getInputOutputRate, makeSwapErrorRight } from '../../../utils'
-import { SOLANA_RANDOM_ADDRESS } from '../utils/constants'
-import { getJupiterPrice, isSupportedChainId } from '../utils/helpers'
+import { referralIdl } from '../idls/referral'
+import {
+  JUPITER_AFFILIATE_CONTRACT_ADDRESS,
+  PDA_ACCOUNT_CREATION_COST,
+  SHAPESHIFT_JUPITER_REFERRAL_KEY,
+  SOLANA_RANDOM_ADDRESS,
+} from '../utils/constants'
+import {
+  getFeeTokenAccountAndInstruction,
+  getJupiterPrice,
+  isSupportedChainId,
+} from '../utils/helpers'
 
 export const getTradeRate = async (
   input: GetTradeRateInput,
@@ -45,6 +57,8 @@ export const getTradeRate = async (
 
   const jupiterUrl = deps.config.REACT_APP_JUPITER_API_URL
 
+  const solAsset = assetsById[solAssetId]
+
   if (!isSupportedChainId(sellAsset.chainId)) {
     return Err(
       makeSwapErrorRight({
@@ -114,6 +128,52 @@ export const getTradeRate = async (
     return { networkFeeCryptoBaseUnit: fast.txFee }
   }
 
+  const buyAssetAddress =
+    buyAsset.assetId === solAssetId
+      ? fromAssetId(wrappedSolAssetId).assetReference
+      : fromAssetId(buyAsset.assetId).assetReference
+
+  const sellAssetAddress =
+    sellAsset.assetId === solAssetId
+      ? fromAssetId(wrappedSolAssetId).assetReference
+      : fromAssetId(sellAsset.assetId).assetReference
+
+  const [buyAssetAccount] = await PublicKey.findProgramAddressSync(
+    [
+      Buffer.from('referral_ata'),
+      new PublicKey(SHAPESHIFT_JUPITER_REFERRAL_KEY).toBuffer(),
+      new PublicKey(buyAssetAddress).toBuffer(),
+    ],
+    new PublicKey(JUPITER_AFFILIATE_CONTRACT_ADDRESS),
+  )
+
+  const [sellAssetAccount] = await PublicKey.findProgramAddressSync(
+    [
+      Buffer.from('referral_ata'),
+      new PublicKey(SHAPESHIFT_JUPITER_REFERRAL_KEY).toBuffer(),
+      new PublicKey(sellAssetAddress).toBuffer(),
+    ],
+    new PublicKey(JUPITER_AFFILIATE_CONTRACT_ADDRESS),
+  )
+
+  const instructionData = new BorshInstructionCoder(referralIdl).encode(
+    'initializeReferralTokenAccount',
+    {},
+  )
+
+  const adapter = deps.assertGetSolanaChainAdapter(sellAsset.chainId)
+
+  const { instruction: feeAccountInstruction } = await getFeeTokenAccountAndInstruction({
+    feePayerPubKey: new PublicKey(input.sendAddress ?? SOLANA_RANDOM_ADDRESS),
+    buyAssetAccount,
+    sellAssetAccount,
+    programId: new PublicKey(JUPITER_AFFILIATE_CONTRACT_ADDRESS),
+    instructionData,
+    buyTokenId: buyAssetAddress,
+    sellTokenId: sellAssetAddress,
+    connection: adapter.getConnection(),
+  })
+
   const protocolFees: Record<AssetId, ProtocolFee> = priceResponse.routePlan.reduce(
     (acc, route) => {
       const feeAssetId = toAssetId({
@@ -139,6 +199,40 @@ export const getTradeRate = async (
     {} as Record<AssetId, ProtocolFee>,
   )
 
+  if (feeAccountInstruction) {
+    if (solAsset) {
+      protocolFees[solAssetId] = {
+        requiresBalance: true,
+        amountCryptoBaseUnit: bnOrZero(protocolFees[solAssetId]?.amountCryptoBaseUnit)
+          .plus(PDA_ACCOUNT_CREATION_COST)
+          .toFixed(),
+        asset: solAsset,
+      }
+    }
+  }
+
+  if (input.sendAddress && solAsset) {
+    const { instruction: createTokenAccountInstruction } = buyAsset
+      ? await adapter.createAssociatedTokenAccountInstruction({
+          from: input.sendAddress,
+          to: receiveAddress ?? input.sendAddress,
+          tokenId: fromAssetId(
+            buyAsset.assetId === solAssetId ? wrappedSolAssetId : buyAsset.assetId,
+          ).assetReference,
+        })
+      : { instruction: undefined }
+
+    if (createTokenAccountInstruction) {
+      protocolFees[solAssetId] = {
+        requiresBalance: true,
+        amountCryptoBaseUnit: bnOrZero(protocolFees[solAssetId]?.amountCryptoBaseUnit)
+          .plus(PDA_ACCOUNT_CREATION_COST)
+          .toFixed(),
+        asset: solAsset,
+      }
+    }
+  }
+
   const rates: TradeRate[] = []
 
   const feeData = await getFeeData()

packages/swapper/src/swappers/JupiterSwapper/utils/constants.ts

@@ -10,6 +10,8 @@ export const JUPITER_SUPPORTED_CHAIN_IDS: SupportedChainIds = {
   buy: jupiterSupportedChainIds,
 }
 
+export const PDA_ACCOUNT_CREATION_COST = 2040000
+
 export const SOLANA_RANDOM_ADDRESS = '2zHKF6tqam3tnNFPK2E9nBDkV7GMXnvdJautmzqQdn8A'
 
 // Jupiter use 40% as a compute unit margin while calculating them, some TX reverts without this

packages/swapper/src/swappers/JupiterSwapper/utils/helpers.ts

@@ -83,19 +83,21 @@ export const getFeeTokenAccountAndInstruction = async ({
   programId,
   buyAssetAccount,
   sellAssetAccount,
-  mint,
+  buyTokenId,
+  sellTokenId,
   instructionData,
   connection,
 }: {
   feePayerPubKey: PublicKey
   programId: PublicKey
   buyAssetAccount: PublicKey
   sellAssetAccount: PublicKey
-  mint: string
+  buyTokenId: string
+  sellTokenId: string
   instructionData: Buffer
   connection: Connection
 }): Promise<{
-  tokenAccount: PublicKey
+  tokenAccount?: PublicKey
   instruction?: TransactionInstruction | undefined
 }> => {
   const buyAssetTokenAccount = await connection.getAccountInfo(buyAssetAccount)
@@ -106,6 +108,16 @@ export const getFeeTokenAccountAndInstruction = async ({
 
   if (sellAssetTokenAccount) return { tokenAccount: sellAssetAccount }
 
+  const buyTokenInfo = await connection.getAccountInfo(new PublicKey(buyTokenId))
+  const sellTokenInfo = await connection.getAccountInfo(new PublicKey(sellTokenId))
+
+  if (
+    buyTokenInfo?.owner.toString() !== TOKEN_PROGRAM_ID.toString() &&
+    sellTokenInfo?.owner.toString() !== TOKEN_PROGRAM_ID.toString()
+  ) {
+    return { tokenAccount: undefined, instruction: undefined }
+  }
+
   const project = new PublicKey(JUPITER_REFERALL_FEE_PROJECT_ACCOUNT)
 
   return {
@@ -133,7 +145,7 @@ export const getFeeTokenAccountAndInstruction = async ({
           isWritable: true,
         },
         {
-          pubkey: new PublicKey(mint),
+          pubkey: new PublicKey(buyTokenId),
           isSigner: false,
           isWritable: false,
         },

src/state/apis/swapper/helpers/validateTradeQuote.ts

@@ -216,6 +216,14 @@ export const validateTradeQuote = (
             const accountId =
               portfolioAccountIdByNumberByChainId[sellAssetAccountNumber][protocolFee.asset.chainId]
             const balanceCryptoBaseUnit = portfolioAccountBalancesBaseUnit[accountId][assetId]
+
+            if (firstHopSellFeeAsset?.assetId === assetId) {
+              return bnOrZero(balanceCryptoBaseUnit)
+                .minus(sellAmountCryptoBaseUnit)
+                .minus(protocolFee.amountCryptoBaseUnit)
+                .lt(0)
+            }
+
             return bnOrZero(balanceCryptoBaseUnit).lt(protocolFee.amountCryptoBaseUnit)
           })
           .map(([_assetId, protocolFee]: [AssetId, ProtocolFee]) => {

yarn.lock

@@ -11699,6 +11699,7 @@ __metadata:
   version: 0.0.0-use.local
   resolution: "@shapeshiftoss/swapper@workspace:packages/swapper"
   dependencies:
+    "@coral-xyz/anchor": 0.29.0
     "@shapeshiftoss/caip": "workspace:^"
     "@shapeshiftoss/chain-adapters": "workspace:^"
     "@shapeshiftoss/types": "workspace:^"
@@ -11774,7 +11775,6 @@ __metadata:
     "@chakra-ui/tag": ^3.1.1
     "@commitlint/cli": ^15.0.0
     "@commitlint/config-conventional": ^15.0.0
-    "@coral-xyz/anchor": 0.29.0
     "@cowprotocol/app-data": ^2.3.0
     "@emotion/react": ^11.13.0
     "@emotion/styled": ^11.13.0