Spike [handlers]: explicit relationship between query and buildCallData signatures

lib/modules/pool/actions/add-liquidity/add-liquidity.types.ts

@@ -1,23 +1,43 @@
-import { AddLiquidityQueryOutput, PriceImpact, TokenAmount } from '@balancer/sdk'
+import { AddLiquidityQueryOutput, TokenAmount } from '@balancer/sdk'
 import { Address } from 'wagmi'
 import { HumanAmountIn } from '../liquidity-types'
+import { UnbalancedAddLiquidityHandler } from './handlers/UnbalancedAddLiquidity.handler'
+import { TwammAddLiquidityHandler } from './handlers/TwammAddLiquidity.handler'
 
-// TODO: this type should be exposed by the SDK
-export type PriceImpactAmount = Awaited<ReturnType<typeof PriceImpact.addLiquidityUnbalanced>>
+export type SupportedHandler = UnbalancedAddLiquidityHandler | TwammAddLiquidityHandler
 
-export type AddLiquidityInputs = {
-  humanAmountsIn: HumanAmountIn[]
-  account?: Address
-  slippagePercent?: string
+export type SdkQueryAddLiquidityOutput = {
+  bptOut: TokenAmount
+  sdkQueryOutput: AddLiquidityQueryOutput
+}
+
+export type TwammQueryAddLiquidityOutput = {
+  bptOut: TokenAmount
 }
 
-// sdkQueryOutput is optional because it will be only used in cases where we use the SDK to query/build the transaction
-// We will probably need a more abstract interface to be used by edge cases
-export type AddLiquidityOutputs = {
+// Default handlers (UnbalancedAddLiquidityHandler) that use the SDK to query/build the transaction will return sdkQueryOutput
+// Edge-case handlers (TwammAddLiquidityHandler) that do not use the SDK edge-case handlers will not return sdkQueryOutput
+export type QueryAddLiquidityOutput<Handler extends SupportedHandler> =
+  Handler extends UnbalancedAddLiquidityHandler
+    ? SdkQueryAddLiquidityOutput
+    : TwammQueryAddLiquidityOutput
+
+export type SdkBuildAddLiquidityInputs = {
+  humanAmountsIn: HumanAmountIn[]
+  account: Address
+  slippagePercent: string
   bptOut: TokenAmount
-  sdkQueryOutput?: AddLiquidityQueryOutput
+  sdkQueryOutput: AddLiquidityQueryOutput
 }
 
-export type BuildLiquidityInputs = {
-  inputs: AddLiquidityInputs
+export type TwammBuildAddLiquidityInputs = {
+  humanAmountsIn: HumanAmountIn[]
+  account: Address
+  slippagePercent: string
+  bptOut: TokenAmount
 }
+
+export type BuildAddLiquidityInputs<Handler extends SupportedHandler> =
+  Handler extends UnbalancedAddLiquidityHandler
+    ? SdkBuildAddLiquidityInputs
+    : TwammBuildAddLiquidityInputs

lib/modules/pool/actions/add-liquidity/handlers/AddLiquidity.handler.ts

@@ -1,21 +1,24 @@
 import { TransactionConfig } from '@/lib/modules/web3/contracts/contract.types'
+import { HumanAmountIn } from '../../liquidity-types'
 import {
-  AddLiquidityInputs,
-  AddLiquidityOutputs,
-  BuildLiquidityInputs,
+  BuildAddLiquidityInputs,
+  QueryAddLiquidityOutput,
+  SupportedHandler,
 } from '../add-liquidity.types'
 
 /**
  * AddLiquidityHandler is an interface that defines the methods that must be implemented by a handler.
  * They take standard inputs from the UI and return frontend standardised outputs.
- * The outputs should not be return types from the SDK. This is to
- * allow handlers to be developed in the future that may not use the SDK.
+ *
+ * The output type of the "query method" and the input type of the "build call data" are generic:
+ * - Default handler types will interact with the SDK to query and build the call data for the transaction
+ * - Edge case handlers (e.g. Twamm handler) will not interact with the SDK.
  */
-export interface AddLiquidityHandler {
-  // Query the SDK for the expected output of adding liquidity
-  queryAddLiquidity(inputs: AddLiquidityInputs): Promise<AddLiquidityOutputs>
+export interface AddLiquidityHandler<Handler extends SupportedHandler> {
+  // Query the expected output of adding liquidity
+  queryAddLiquidity(humanAmountsIn: HumanAmountIn[]): Promise<QueryAddLiquidityOutput<Handler>>
   // Calculate the price impact of adding liquidity
-  calculatePriceImpact(inputs: AddLiquidityInputs): Promise<number>
+  calculatePriceImpact(humanAmountsIn: HumanAmountIn[]): Promise<number>
   // Build tx payload for adding liquidity
-  buildAddLiquidityCallData(inputs: BuildLiquidityInputs): Promise<TransactionConfig>
+  buildAddLiquidityCallData(inputs: BuildAddLiquidityInputs<Handler>): Promise<TransactionConfig>
 }

lib/modules/pool/actions/add-liquidity/handlers/TwammAddLiquidity.handler.ts

@@ -2,25 +2,22 @@ import { SupportedChainId } from '@/lib/config/config.types'
 import { TransactionConfig } from '@/lib/modules/web3/contracts/contract.types'
 import { emptyAddress } from '@/lib/modules/web3/contracts/wagmi-helpers'
 import { Token, TokenAmount } from '@balancer/sdk'
-import {
-  AddLiquidityInputs,
-  AddLiquidityOutputs,
-  BuildLiquidityInputs,
-} from '../add-liquidity.types'
+import { HumanAmountIn } from '../../liquidity-types'
+import { TwammBuildAddLiquidityInputs, TwammQueryAddLiquidityOutput } from '../add-liquidity.types'
 import { AddLiquidityHandler } from './AddLiquidity.handler'
 
 /**
  * TwammAddLiquidityHandler is a handler that implements the
  * AddLiquidityHandler interface for TWAMM adds.
  * This is just a fake example to show how to implement edge-case handlers.
  */
-export class TwammAddLiquidityHandler implements AddLiquidityHandler {
+export class TwammAddLiquidityHandler implements AddLiquidityHandler<TwammAddLiquidityHandler> {
   constructor(private chainId: SupportedChainId) {}
 
   // TODO: This is a non-sense example implementation
-  public async queryAddLiquidity({
-    humanAmountsIn,
-  }: AddLiquidityInputs): Promise<AddLiquidityOutputs> {
+  public async queryAddLiquidity(
+    humanAmountsIn: HumanAmountIn[]
+  ): Promise<TwammQueryAddLiquidityOutput> {
     const tokenAmount = TokenAmount.fromHumanAmount(
       {} as unknown as Token,
       humanAmountsIn[0].humanAmount || '0'
@@ -30,15 +27,16 @@ export class TwammAddLiquidityHandler implements AddLiquidityHandler {
   }
 
   // TODO: This is a non-sense example implementation
-  public async calculatePriceImpact({ humanAmountsIn }: AddLiquidityInputs): Promise<number> {
+  public async calculatePriceImpact(humanAmountsIn: HumanAmountIn[]): Promise<number> {
     return Number(humanAmountsIn[0].humanAmount)
   }
 
   // TODO: This is a non-sense example implementation
-  public async buildAddLiquidityCallData(
-    buildInputs: BuildLiquidityInputs
-  ): Promise<TransactionConfig> {
-    const { humanAmountsIn, account, slippagePercent } = buildInputs.inputs
+  public async buildAddLiquidityCallData({
+    humanAmountsIn,
+    account,
+    slippagePercent,
+  }: TwammBuildAddLiquidityInputs): Promise<TransactionConfig> {
     if (!account || !slippagePercent) throw new Error('Missing account or slippage')
 
     const value = BigInt(humanAmountsIn[0].humanAmount)

lib/modules/pool/actions/add-liquidity/handlers/UnbalancedAddLiquidity.handler.integration.spec.ts

@@ -6,12 +6,13 @@ import { HumanAmount } from '@balancer/sdk'
 import { Address } from 'viem'
 import { aPhantomStablePoolStateInputMock } from '../../../__mocks__/pool.builders'
 import { Pool } from '../../../usePool'
-import { selectAddLiquidityHandler } from './selectAddLiquidityHandler'
 import { HumanAmountIn } from '../../liquidity-types'
+import { SdkBuildAddLiquidityInputs } from '../add-liquidity.types'
+import { UnbalancedAddLiquidityHandler } from './UnbalancedAddLiquidity.handler'
+import { selectAddLiquidityHandler } from './selectAddLiquidityHandler'
 
 function selectUnbalancedHandler() {
-  //TODO: refactor mock builders to build poolStateInput and pool at the same time
-  return selectAddLiquidityHandler(aWjAuraWethPoolElementMock())
+  return selectAddLiquidityHandler(aWjAuraWethPoolElementMock()) as UnbalancedAddLiquidityHandler
 }
 
 describe('When adding unbalanced liquidity for a weighted  pool', () => {
@@ -23,7 +24,7 @@ describe('When adding unbalanced liquidity for a weighted  pool', () => {
       { humanAmount: '1', tokenAddress: wjAuraAddress },
     ]
 
-    const priceImpact = await handler.calculatePriceImpact({ humanAmountsIn })
+    const priceImpact = await handler.calculatePriceImpact(humanAmountsIn)
     expect(priceImpact).toBeGreaterThan(0.002)
   })
 
@@ -35,7 +36,7 @@ describe('When adding unbalanced liquidity for a weighted  pool', () => {
       { humanAmount: '', tokenAddress: balAddress },
     ]
 
-    const priceImpact = await handler.calculatePriceImpact({ humanAmountsIn })
+    const priceImpact = await handler.calculatePriceImpact(humanAmountsIn)
 
     expect(priceImpact).toEqual(0)
   })
@@ -48,9 +49,7 @@ describe('When adding unbalanced liquidity for a weighted  pool', () => {
 
     const handler = selectUnbalancedHandler()
 
-    const result = await handler.queryAddLiquidity({
-      humanAmountsIn,
-    })
+    const result = await handler.queryAddLiquidity(humanAmountsIn)
 
     expect(result.bptOut.amount).toBeGreaterThan(300000000000000000000n)
   })
@@ -63,16 +62,16 @@ describe('When adding unbalanced liquidity for a weighted  pool', () => {
 
     const handler = selectUnbalancedHandler()
 
-    await handler.queryAddLiquidity({
-      humanAmountsIn,
-    })
+    const { sdkQueryOutput } = await handler.queryAddLiquidity(humanAmountsIn)
 
-    const inputs = {
+    const buildInputs: SdkBuildAddLiquidityInputs = {
       humanAmountsIn,
       account: defaultTestUserAccount,
       slippagePercent: '0.2',
+      bptOut: sdkQueryOutput.bptOut,
+      sdkQueryOutput,
     }
-    const result = await handler.buildAddLiquidityCallData({ inputs })
+    const result = await handler.buildAddLiquidityCallData(buildInputs)
 
     expect(result.to).toBe(networkConfig.contracts.balancer.vaultV2)
     expect(result.data).toBeDefined()
@@ -94,7 +93,7 @@ describe('When adding unbalanced liquidity for a stable pool', () => {
       }
     })
 
-    const priceImpact = await handler.calculatePriceImpact({ humanAmountsIn })
+    const priceImpact = await handler.calculatePriceImpact(humanAmountsIn)
     expect(priceImpact).toBeGreaterThan(0.001)
   })
 })

lib/modules/pool/actions/add-liquidity/handlers/UnbalancedAddLiquidity.handler.ts

@@ -1,23 +1,19 @@
+/* eslint-disable @typescript-eslint/no-non-null-assertion */
 import { getDefaultRpcUrl } from '@/lib/modules/web3/Web3Provider'
 import { TransactionConfig } from '@/lib/modules/web3/contracts/contract.types'
 import {
   AddLiquidity,
   AddLiquidityKind,
-  AddLiquidityQueryOutput,
   AddLiquidityUnbalancedInput,
   PriceImpact,
+  PriceImpactAmount,
   Slippage,
 } from '@balancer/sdk'
-import {
-  AddLiquidityInputs,
-  AddLiquidityOutputs,
-  BuildLiquidityInputs,
-  PriceImpactAmount,
-} from '../add-liquidity.types'
-import { AddLiquidityHandler } from './AddLiquidity.handler'
 import { Pool } from '../../../usePool'
 import { LiquidityActionHelpers, areEmptyAmounts } from '../../LiquidityActionHelpers'
 import { HumanAmountIn } from '../../liquidity-types'
+import { SdkBuildAddLiquidityInputs, SdkQueryAddLiquidityOutput } from '../add-liquidity.types'
+import { AddLiquidityHandler } from './AddLiquidity.handler'
 
 /**
  * UnbalancedAddLiquidityHandler is a handler that implements the
@@ -26,26 +22,27 @@ import { HumanAmountIn } from '../../liquidity-types'
  * methods. It also handles the case where one of the input tokens is the native
  * asset instead of the wrapped native asset.
  */
-export class UnbalancedAddLiquidityHandler implements AddLiquidityHandler {
+export class UnbalancedAddLiquidityHandler
+  implements AddLiquidityHandler<UnbalancedAddLiquidityHandler>
+{
   helpers: LiquidityActionHelpers
-  sdkQueryOutput?: AddLiquidityQueryOutput
 
   constructor(pool: Pool) {
     this.helpers = new LiquidityActionHelpers(pool)
   }
 
-  public async queryAddLiquidity({
-    humanAmountsIn,
-  }: AddLiquidityInputs): Promise<AddLiquidityOutputs> {
+  public async queryAddLiquidity(
+    humanAmountsIn: HumanAmountIn[]
+  ): Promise<SdkQueryAddLiquidityOutput> {
     const addLiquidity = new AddLiquidity()
     const addLiquidityInput = this.constructSdkInput(humanAmountsIn)
 
-    this.sdkQueryOutput = await addLiquidity.query(addLiquidityInput, this.helpers.poolStateInput)
+    const sdkQueryOutput = await addLiquidity.query(addLiquidityInput, this.helpers.poolStateInput)
 
-    return { bptOut: this.sdkQueryOutput.bptOut }
+    return { bptOut: sdkQueryOutput.bptOut, sdkQueryOutput }
   }
 
-  public async calculatePriceImpact({ humanAmountsIn }: AddLiquidityInputs): Promise<number> {
+  public async calculatePriceImpact(humanAmountsIn: HumanAmountIn[]): Promise<number> {
     if (areEmptyAmounts(humanAmountsIn)) {
       // Avoid price impact calculation when there are no amounts in
       return 0
@@ -64,23 +61,15 @@ export class UnbalancedAddLiquidityHandler implements AddLiquidityHandler {
   /*
     sdkQueryOutput is the result of the query that we run in the add liquidity form
   */
-  public async buildAddLiquidityCallData(
-    buildInputs: BuildLiquidityInputs
-  ): Promise<TransactionConfig> {
-    const { account, slippagePercent } = buildInputs.inputs
-    if (!account || !slippagePercent) throw new Error('Missing account or slippage')
-    if (!this.sdkQueryOutput) {
-      console.error('Missing sdkQueryOutput.')
-      throw new Error(
-        `Missing sdkQueryOutput.
-It looks that you did not call useAddLiquidityBtpOutQuery before trying to build the tx config`
-      )
-    }
-
+  public async buildAddLiquidityCallData({
+    account,
+    slippagePercent,
+    sdkQueryOutput,
+  }: SdkBuildAddLiquidityInputs): Promise<TransactionConfig> {
     const addLiquidity = new AddLiquidity()
 
     const { call, to, value } = addLiquidity.buildCall({
-      ...this.sdkQueryOutput,
+      ...sdkQueryOutput,
       slippage: Slippage.fromPercentage(`${Number(slippagePercent)}`),
       sender: account,
       recipient: account,

lib/modules/pool/actions/add-liquidity/handlers/selectAddLiquidityHandler.ts

@@ -1,12 +1,12 @@
 import { getChainId } from '@/lib/config/app.config'
 import { Pool } from '../../../usePool'
-import { AddLiquidityHandler } from './AddLiquidity.handler'
+import { SupportedHandler } from '../add-liquidity.types'
 import { TwammAddLiquidityHandler } from './TwammAddLiquidity.handler'
 import { UnbalancedAddLiquidityHandler } from './UnbalancedAddLiquidity.handler'
 
 export function selectAddLiquidityHandler(pool: Pool) {
   // TODO: Depending on the pool attributes we will return a different handler
-  let handler: AddLiquidityHandler
+  let handler: SupportedHandler
   if (pool.id === 'TWAMM-example') {
     // This is just an example to illustrate how edge-case handlers would receive different inputs but return a common contract
     handler = new TwammAddLiquidityHandler(getChainId(pool.chain))

lib/modules/pool/actions/add-liquidity/queries/useAddLiquidityBuildCallDataQuery.ts

@@ -6,15 +6,22 @@ import { useUserAccount } from '@/lib/modules/web3/useUserAccount'
 import { useQuery } from 'wagmi'
 import { Pool } from '../../../usePool'
 import { HumanAmountIn } from '../../liquidity-types'
-import { AddLiquidityHandler } from '../handlers/AddLiquidity.handler'
+import {
+  BuildAddLiquidityInputs,
+  QueryAddLiquidityOutput,
+  SdkBuildAddLiquidityInputs,
+  SupportedHandler,
+} from '../add-liquidity.types'
+import { TwammAddLiquidityHandler } from '../handlers/TwammAddLiquidity.handler'
 import { addLiquidityKeys } from './add-liquidity-keys'
 
 // Uses the SDK to build a transaction config to be used by wagmi's useManagedSendTransaction
 export function useAddLiquidityBuildCallDataQuery(
-  handler: AddLiquidityHandler,
+  handler: SupportedHandler,
   humanAmountsIn: HumanAmountIn[],
   isActiveStep: boolean,
-  pool: Pool
+  pool: Pool,
+  queryOutput: QueryAddLiquidityOutput<SupportedHandler>
 ) {
   const { userAddress, isConnected } = useUserAccount()
   const { slippage } = useUserSettings()
@@ -29,16 +36,30 @@ export function useAddLiquidityBuildCallDataQuery(
       humanAmountsIn,
     }),
     async () => {
-      const inputs = {
+      const baseInput: BuildAddLiquidityInputs<SupportedHandler> = {
         humanAmountsIn,
         account: userAddress,
         slippagePercent: slippage,
+        bptOut: queryOutput.bptOut,
+      }
+
+      const isSdkHandler = 'sdkQueryOutput' in queryOutput && queryOutput.sdkQueryOutput
+
+      if (isSdkHandler) {
+        const sdkBuildInput: SdkBuildAddLiquidityInputs = {
+          ...baseInput,
+          sdkQueryOutput: queryOutput.sdkQueryOutput,
+        }
+        return handler.buildAddLiquidityCallData(sdkBuildInput)
+      }
+      if (handler instanceof TwammAddLiquidityHandler) {
+        return handler.buildAddLiquidityCallData(baseInput)
       }
-      return handler.buildAddLiquidityCallData({ inputs })
     },
     {
       enabled:
         isActiveStep && // If the step is not active (the user did not click Next button) avoid running the build tx query to save RPC requests
+        queryOutput.bptOut && // undefined bptOut means that the preview query did not finish yet
         isConnected &&
         hasAllowances(humanAmountsIn, pool),
     }

lib/modules/pool/actions/add-liquidity/queries/useAddLiquidityPreviewQuery.integration.spec.tsx

@@ -24,7 +24,7 @@ test('queries btp out for add liquidity', async () => {
 
   const result = await testQuery(humanAmountsIn)
 
-  await waitFor(() => expect(result.current.bptOut).not.toBeNull())
+  await waitFor(() => expect(result.current.bptOut).toBeDefined())
 
   expect(result.current.bptOut?.amount).toBeDefined()
   expect(result.current.isPreviewQueryLoading).toBeFalsy()