[PE-6141] Add swap functionality and swap UI for Wallet update M2

packages/common/src/api/index.ts

@@ -130,6 +130,8 @@ export * from './tan-query/wallets/useTokenPrice'
 export * from './tan-query/wallets/useWalletCollectibles'
 export * from './tan-query/wallets/useWalletOwner'
 export * from './tan-query/wallets/useUSDCBalance'
+export * from './tan-query/jupiter/useSwapTokens'
+export * from './tan-query/jupiter/useTokenExchangeRate'
 
 // Saga fetch utils, remove when migration is complete
 export * from './tan-query/saga-utils'

packages/common/src/api/tan-query/jupiter/constants.ts

@@ -0,0 +1,25 @@
+import { PublicKey } from '@solana/web3.js'
+
+import { TOKEN_LISTING_MAP } from '~/store/ui/buy-audio/constants'
+
+import { UserBankManagedTokenInfo } from './types'
+
+export const JUPITER_PROGRAM_ID = new PublicKey(
+  'JUP6LkbZbjS1jKKwapdHNy74zcZ3tLUZoi5QNyVTaV4'
+)
+
+export const USER_BANK_MANAGED_TOKENS: Record<
+  string,
+  UserBankManagedTokenInfo
+> = {
+  [TOKEN_LISTING_MAP.AUDIO.address.toUpperCase()]: {
+    mintAddress: TOKEN_LISTING_MAP.AUDIO.address,
+    claimableTokenMint: 'wAUDIO',
+    decimals: TOKEN_LISTING_MAP.AUDIO.decimals
+  },
+  [TOKEN_LISTING_MAP.USDC.address.toUpperCase()]: {
+    mintAddress: TOKEN_LISTING_MAP.USDC.address,
+    claimableTokenMint: 'USDC',
+    decimals: TOKEN_LISTING_MAP.USDC.decimals
+  }
+}

packages/common/src/api/tan-query/jupiter/types.ts

@@ -0,0 +1,48 @@
+export enum SwapStatus {
+  IDLE = 'IDLE',
+  GETTING_QUOTE = 'GETTING_QUOTE',
+  BUILDING_TRANSACTION = 'BUILDING_TRANSACTION',
+  SENDING_TRANSACTION = 'SENDING_TRANSACTION',
+  SUCCESS = 'SUCCESS',
+  ERROR = 'ERROR'
+}
+
+export enum SwapErrorType {
+  WALLET_ERROR = 'WALLET_ERROR',
+  QUOTE_FAILED = 'QUOTE_FAILED',
+  BUILD_FAILED = 'BUILD_FAILED',
+  RELAY_FAILED = 'RELAY_FAILED',
+  SIMULATION_FAILED = 'SIMULATION_FAILED',
+  UNKNOWN = 'UNKNOWN'
+}
+
+export type SwapTokensParams = {
+  inputMint: string
+  outputMint: string
+  amountUi: number
+  slippageBps?: number
+  wrapUnwrapSol?: boolean
+}
+
+export type SwapTokensResult = {
+  status: SwapStatus
+  signature?: string
+  error?: {
+    type: SwapErrorType
+    message: string
+  }
+  inputAmount?: {
+    amount: number
+    uiAmount: number
+  }
+  outputAmount?: {
+    amount: number
+    uiAmount: number
+  }
+}
+
+export interface UserBankManagedTokenInfo {
+  mintAddress: string
+  claimableTokenMint: 'wAUDIO' | 'USDC'
+  decimals: number
+}

packages/common/src/api/tan-query/jupiter/useSwapTokens.ts

@@ -0,0 +1,211 @@
+import { createCloseAccountInstruction } from '@solana/spl-token'
+import {
+  PublicKey,
+  TransactionInstruction,
+  VersionedTransaction
+} from '@solana/web3.js'
+import { useMutation, useQueryClient } from '@tanstack/react-query'
+
+import { useGetCurrentUser } from '~/api'
+import { useAudiusQueryContext } from '~/audius-query'
+import { Feature } from '~/models'
+import {
+  convertJupiterInstructions,
+  getJupiterQuoteByMint,
+  jupiterInstance
+} from '~/services/Jupiter'
+
+import { QUERY_KEYS } from '../queryKeys'
+
+import { USER_BANK_MANAGED_TOKENS } from './constants'
+import {
+  SwapErrorType,
+  SwapStatus,
+  SwapTokensParams,
+  SwapTokensResult
+} from './types'
+import { addUserBankToAtaInstructions, getSwapErrorResponse } from './utils'
+
+/**
+ * Hook for executing token swaps using Jupiter.
+ * Swaps any supported SPL token (or SOL) for another.
+ */
+export const useSwapTokens = () => {
+  const queryClient = useQueryClient()
+  const { solanaWalletService, reportToSentry, audiusSdk } =
+    useAudiusQueryContext()
+  const { data: user } = useGetCurrentUser({})
+
+  return useMutation<SwapTokensResult, Error, SwapTokensParams>({
+    mutationFn: async (params): Promise<SwapTokensResult> => {
+      const {
+        inputMint: inputMintUiAddress,
+        outputMint: outputMintUiAddress,
+        amountUi
+      } = params
+      const slippageBps = params.slippageBps ?? 50
+      const wrapUnwrapSol = params.wrapUnwrapSol ?? true
+
+      let quoteResult
+      let signature: string | undefined
+      let errorStage = 'UNKNOWN'
+      const instructions: TransactionInstruction[] = []
+
+      try {
+        // ---------- 1. Initialize Dependencies & Wallet ----------
+        errorStage = 'WALLET_INITIALIZATION'
+        const [sdk, keypair] = await Promise.all([
+          audiusSdk(),
+          solanaWalletService.getKeypair()
+        ])
+
+        if (!keypair) {
+          return {
+            status: SwapStatus.ERROR,
+            error: {
+              type: SwapErrorType.WALLET_ERROR,
+              message: 'Wallet not initialised'
+            }
+          }
+        }
+        const userPublicKey = keypair.publicKey
+        const feePayer = await sdk.services.solanaClient.getFeePayer()
+        const ethAddress = user?.wallet
+
+        // ---------- 2. Get Quote from Jupiter ----------
+        errorStage = 'QUOTE_RETRIEVAL'
+        quoteResult = await getJupiterQuoteByMint({
+          inputMint: inputMintUiAddress,
+          outputMint: outputMintUiAddress,
+          amountUi,
+          slippageBps,
+          swapMode: 'ExactIn',
+          onlyDirectRoutes: true
+        })
+
+        // ---------- 3. Prepare Transaction Instructions ----------
+        const inputTokenConfig =
+          USER_BANK_MANAGED_TOKENS[inputMintUiAddress.toUpperCase()]
+        const outputTokenConfig =
+          USER_BANK_MANAGED_TOKENS[outputMintUiAddress.toUpperCase()]
+
+        // --- 3a. Handle Input Token (if user bank managed) ---
+        errorStage = 'INPUT_TOKEN_PREPARATION'
+
+        const sourceAtaForJupiter = await addUserBankToAtaInstructions({
+          tokenInfo: inputTokenConfig!,
+          userPublicKey,
+          ethAddress: ethAddress!,
+          amountLamports: BigInt(quoteResult.inputAmount.amount),
+          sdk,
+          feePayer,
+          instructions
+        })
+
+        // --- 3b. Determine Jupiter's Destination Token Account ---
+        errorStage = 'OUTPUT_TOKEN_PREPARATION'
+
+        const result =
+          await sdk.services.claimableTokensClient.getOrCreateUserBank({
+            ethWallet: ethAddress!,
+            mint: outputTokenConfig!.claimableTokenMint
+          })
+        const outputUserBankAddress = result.userBank
+        const preferredJupiterDestination = outputUserBankAddress.toBase58()
+
+        // --- 3c. Get Jupiter Swap Instructions ---
+        errorStage = 'SWAP_INSTRUCTION_RETRIEVAL'
+        const {
+          tokenLedgerInstruction,
+          swapInstruction,
+          addressLookupTableAddresses
+        } = await jupiterInstance.swapInstructionsPost({
+          swapRequest: {
+            quoteResponse: quoteResult.quote,
+            userPublicKey: userPublicKey.toBase58(),
+            destinationTokenAccount: preferredJupiterDestination,
+            wrapAndUnwrapSol: wrapUnwrapSol,
+            useSharedAccounts: true
+          }
+        })
+
+        const jupiterInstructions = convertJupiterInstructions([
+          tokenLedgerInstruction,
+          swapInstruction
+        ])
+
+        instructions.push(...jupiterInstructions)
+
+        // --- 3f. Add Cleanup Instructions for Temporary ATAs ---
+        const atasToClose: PublicKey[] = []
+
+        // Add source ATA if it's temporary
+        if (sourceAtaForJupiter) {
+          atasToClose.push(sourceAtaForJupiter)
+        }
+
+        // Add close account instructions for all ATAs that need to be closed
+        for (const ataToClose of atasToClose) {
+          instructions.push(
+            createCloseAccountInstruction(ataToClose, feePayer, userPublicKey)
+          )
+        }
+
+        // ---------- 4. Build and Sign Transaction ----------
+        errorStage = 'TRANSACTION_BUILD'
+        const swapTx: VersionedTransaction =
+          await sdk.services.solanaClient.buildTransaction({
+            feePayer,
+            instructions,
+            addressLookupTables: addressLookupTableAddresses.map(
+              (addr: string) => new PublicKey(addr)
+            )
+          })
+
+        swapTx.sign([keypair])
+
+        // ---------- 5. Send Transaction ----------
+        errorStage = 'TRANSACTION_RELAY'
+        signature = await sdk.services.solanaClient.sendTransaction(swapTx)
+
+        // ---------- 6. Success & Invalidation ----------
+        if (user?.wallet) {
+          queryClient.invalidateQueries({
+            queryKey: [QUERY_KEYS.usdcBalance, user.wallet]
+          })
+          queryClient.invalidateQueries({
+            queryKey: [QUERY_KEYS.audioBalance]
+          })
+        }
+
+        return {
+          status: SwapStatus.SUCCESS,
+          signature,
+          inputAmount: quoteResult.inputAmount,
+          outputAmount: quoteResult.outputAmount
+        }
+      } catch (error: unknown) {
+        reportToSentry({
+          name: `JupiterSwap${errorStage}Error`,
+          error: error as Error,
+          feature: Feature.TanQuery,
+          additionalInfo: {
+            params,
+            signature,
+            quoteResponse: quoteResult?.quote
+          }
+        })
+
+        return getSwapErrorResponse({
+          errorStage,
+          error: error as Error,
+          inputAmount: quoteResult?.inputAmount,
+          outputAmount: quoteResult?.outputAmount
+        })
+      }
+    },
+    onMutate: () => {
+      return { status: SwapStatus.SENDING_TRANSACTION }
+    }
+  })
+}