Fix token selection bugs with lockChainId and restrictedTokens

.changeset/brown-cougars-exist.md

@@ -0,0 +1,5 @@
+---
+'@reservoir0x/relay-kit-ui': patch
+---
+
+Fix token selection bugs with lockChainId and restrictedTokens

packages/ui/src/components/common/TokenSelector/TokenSelector.tsx

@@ -149,14 +149,27 @@ const TokenSelector: FC<TokenSelectorProps> = ({
         )
       : configuredChains
 
+  const filteredRestrictedTokenList = useMemo(() => {
+    if (!restrictedTokensList) return undefined
+
+    // Only filter tokens if we're viewing a specific chain
+    if (chainFilter.id === undefined) {
+      return restrictedTokensList
+    }
+
+    return restrictedTokensList.filter(
+      (token) => token.chainId === chainFilter.id
+    )
+  }, [restrictedTokensList, chainFilter.id])
+
   const useDefaultTokenList =
     debouncedTokenSearchValue === '' &&
-    (!restrictedTokensList || !restrictedTokensList.length)
+    (!filteredRestrictedTokenList || !filteredRestrictedTokenList.length)
 
   let tokenListQuery: string[] | undefined
 
-  if (restrictedTokensList && restrictedTokensList.length > 0) {
-    tokenListQuery = restrictedTokensList.map(
+  if (filteredRestrictedTokenList && filteredRestrictedTokenList.length > 0) {
+    tokenListQuery = filteredRestrictedTokenList.map(
       (token) => `${token.chainId}:${token.address}`
     )
   }
@@ -217,7 +230,7 @@ const TokenSelector: FC<TokenSelectorProps> = ({
       : undefined
   )
 
-  const restrictedTokenAddresses = restrictedTokensList?.map((token) =>
+  const restrictedTokenAddresses = filteredRestrictedTokenList?.map((token) =>
     token.address.toLowerCase()
   )
 
@@ -234,8 +247,8 @@ const TokenSelector: FC<TokenSelectorProps> = ({
 
   let suggestedTokenQuery: string[] | undefined
 
-  if (restrictedTokensList && restrictedTokensList.length > 0) {
-    suggestedTokenQuery = restrictedTokensList.map(
+  if (filteredRestrictedTokenList && filteredRestrictedTokenList.length > 0) {
+    suggestedTokenQuery = filteredRestrictedTokenList.map(
       (token) => `${token.chainId}:${token.address}`
     )
   } else if (duneTokenBalances) {
@@ -577,6 +590,7 @@ const TokenSelector: FC<TokenSelectorProps> = ({
                 setInputElement={setInputElement}
                 tokenSearchInput={tokenSearchInput}
                 setTokenSearchInput={setTokenSearchInput}
+                configuredChainIds={configuredChainIds}
                 chainFilterOptions={chainFilterOptions}
                 chainFilter={chainFilter}
                 setChainFilter={setChainFilter}

packages/ui/src/components/common/TokenSelector/steps/SetCurrencyStep.tsx

@@ -44,6 +44,7 @@ type SetCurrencyProps = {
   ) => void
   tokenSearchInput: string
   setTokenSearchInput: (value: string) => void
+  configuredChainIds: number[] | undefined
   chainFilterOptions: RelayChain[]
   chainFilter: ChainFilterValue
   setChainFilter: (value: React.SetStateAction<ChainFilterValue>) => void
@@ -73,6 +74,7 @@ export const SetCurrencyStep: FC<SetCurrencyProps> = ({
   setInputElement,
   tokenSearchInput,
   setTokenSearchInput,
+  configuredChainIds,
   chainFilterOptions,
   chainFilter,
   setChainFilter,
@@ -158,7 +160,7 @@ export const SetCurrencyStep: FC<SetCurrencyProps> = ({
         Select Token
       </Text>
       <Flex css={{ width: '100%', gap: '3', height: '400px' }}>
-        {isDesktop && allChains.length > 2 ? (
+        {isDesktop && (!configuredChainIds || configuredChainIds.length > 1) ? (
           <>
             <Flex
               direction="column"
@@ -384,7 +386,8 @@ export const SetCurrencyStep: FC<SetCurrencyProps> = ({
                   }
                 />
               </AccessibleListItem>
-              {!isDesktop && allChains.length > 2 ? (
+              {!isDesktop &&
+              (!configuredChainIds || configuredChainIds.length > 1) ? (
                 <ChainFilter
                   options={allChains}
                   value={chainFilter}

packages/ui/src/components/widgets/SwapWidget/index.tsx

@@ -36,9 +36,8 @@ import {
 } from '@reservoir0x/relay-sdk'
 import SwapRouteSelector from '../SwapRouteSelector.js'
 import { ProviderOptionsContext } from '../../../providers/RelayKitProvider.js'
-import { faCircleExclamation } from '@fortawesome/free-solid-svg-icons'
-import Tooltip from '../../primitives/Tooltip.js'
 import { findBridgableToken } from '../../../utils/tokens.js'
+import { isChainLocked } from '../../../utils/tokenSelector.js'
 
 type BaseSwapWidgetProps = {
   defaultFromToken?: Token
@@ -385,10 +384,11 @@ const SwapWidget: FC<SwapWidgetProps> = ({
                       <ChainTrigger
                         token={fromToken}
                         chain={fromChain}
-                        locked={
-                          lockChainId !== undefined &&
-                          lockChainId === fromToken?.chainId
-                        }
+                        locked={isChainLocked(
+                          fromToken?.chainId,
+                          lockChainId,
+                          toToken?.chainId
+                        )}
                         onClick={() => {
                           setFromTokenSelectorType('chain')
                           fromTokenSelectorOpenState[1](
@@ -414,8 +414,8 @@ const SwapWidget: FC<SwapWidgetProps> = ({
                           tradeType === 'EXACT_INPUT'
                             ? amountInputValue
                             : amountInputValue
-                              ? formatFixedLength(amountInputValue, 8)
-                              : amountInputValue
+                            ? formatFixedLength(amountInputValue, 8)
+                            : amountInputValue
                         }
                         setValue={(e) => {
                           setAmountInputValue(e)
@@ -475,10 +475,13 @@ const SwapWidget: FC<SwapWidgetProps> = ({
                         lockedChainIds={
                           isSingleChainLocked
                             ? [lockChainId]
-                            : fromToken?.chainId !== undefined &&
-                                fromToken?.chainId === lockChainId
-                              ? [fromToken?.chainId]
-                              : undefined
+                            : isChainLocked(
+                                fromToken?.chainId,
+                                lockChainId,
+                                toToken?.chainId
+                              ) && fromToken?.chainId
+                            ? [fromToken.chainId]
+                            : undefined
                         }
                         chainIdsFilter={
                           !fromChainWalletVMSupported && toToken
@@ -746,10 +749,11 @@ const SwapWidget: FC<SwapWidgetProps> = ({
                       <ChainTrigger
                         token={toToken}
                         chain={toChain}
-                        locked={
-                          lockChainId !== undefined &&
-                          lockChainId === toToken?.chainId
-                        }
+                        locked={isChainLocked(
+                          toToken?.chainId,
+                          lockChainId,
+                          fromToken?.chainId
+                        )}
                         onClick={() => {
                           setToTokenSelectorType('chain')
                           toTokenSelectorOpenState[1](
@@ -775,8 +779,8 @@ const SwapWidget: FC<SwapWidgetProps> = ({
                           tradeType === 'EXPECTED_OUTPUT'
                             ? amountOutputValue
                             : amountOutputValue
-                              ? formatFixedLength(amountOutputValue, 8)
-                              : amountOutputValue
+                            ? formatFixedLength(amountOutputValue, 8)
+                            : amountOutputValue
                         }
                         setValue={(e) => {
                           setAmountOutputValue(e)
@@ -864,10 +868,13 @@ const SwapWidget: FC<SwapWidgetProps> = ({
                         lockedChainIds={
                           isSingleChainLocked
                             ? [lockChainId]
-                            : toToken?.chainId !== undefined &&
-                                toToken?.chainId === lockChainId
-                              ? [toToken?.chainId]
-                              : undefined
+                            : isChainLocked(
+                                toToken?.chainId,
+                                lockChainId,
+                                fromToken?.chainId
+                              ) && toToken?.chainId
+                            ? [toToken.chainId]
+                            : undefined
                         }
                         chainIdsFilter={
                           !fromChainWalletVMSupported && fromToken

packages/ui/src/utils/tokenSelector.ts

@@ -0,0 +1,14 @@
+export const isChainLocked = (
+  chainId: number | undefined,
+  lockChainId: number | undefined,
+  otherTokenChainId: number | undefined
+) => {
+  if (lockChainId === undefined) return false
+
+  // If this token is on the locked chain, only lock it if the other token isn't
+  if (chainId === lockChainId) {
+    return otherTokenChainId !== lockChainId
+  }
+
+  return false
+}