apply review

Author: RogerLamTd

packages/ui/src/components/WalletContainer.tsx

@@ -1,6 +1,8 @@
-import { PUB_CHAIN, PUB_CHAIN_NAME, TAIKO_L2_CHAIN_ID } from "@/constants";
+import { PUB_CHAIN, PUB_CHAIN_NAME } from "@/constants";
+import { useWalletChainPolicy } from "@/context/WalletChainPolicy";
 import { config } from "@/context/Web3Modal";
 import { formatHexString } from "@/utils/evm";
+import { shouldForcePrimaryChainSwitch } from "@/utils/wallet-chain-policy";
 import { MemberAvatar } from "@aragon/ods";
 import { useWeb3Modal } from "@web3modal/wagmi/react";
 import classNames from "classnames";
@@ -16,6 +18,7 @@ const WalletContainer = () => {
   const { open } = useWeb3Modal();
   const { address, isConnected, chainId } = useAccount();
   const { switchChain } = useSwitchChain();
+  const { allowedSecondaryChainIds } = useWalletChainPolicy();
 
   const { data: ensName } = useEnsName({
     config,
@@ -32,11 +35,10 @@ const WalletContainer = () => {
   });
 
   useEffect(() => {
-    if (!chainId) return;
-    else if (chainId === PUB_CHAIN.id || chainId === TAIKO_L2_CHAIN_ID ) return;
+    if (!shouldForcePrimaryChainSwitch(chainId, PUB_CHAIN.id, allowedSecondaryChainIds)) return;
 
     switchChain({ chainId: PUB_CHAIN.id });
-  }, [chainId, switchChain]);
+  }, [allowedSecondaryChainIds, chainId, switchChain]);
 
   return (
     <button

packages/ui/src/components/l2Execution/ProposalL2Execution.tsx

@@ -1,8 +1,10 @@
+import { useEffect } from "react";
 import { AlertInline, Button, Spinner } from "@aragon/ods";
 import { useWeb3Modal } from "@web3modal/wagmi/react";
 import { type Hex } from "viem";
 import { useAccount, useSwitchChain } from "wagmi";
 import { PUB_TAIKO_BRIDGE_ADDRESS, TAIKO_L2_CHAIN_ID } from "@/constants";
+import { useWalletChainPolicy } from "@/context/WalletChainPolicy";
 import { useL2AnchorSync } from "@/hooks/useL2AnchorSync";
 import { useL2LegExecution } from "@/hooks/useL2LegExecution";
 import { shouldRenderL2ExecutionCard } from "@/utils/l2-execution";
@@ -37,6 +39,7 @@ export function ProposalL2Execution({
   const { isConnected, chain } = useAccount();
   const { open } = useWeb3Modal();
   const { switchChain } = useSwitchChain();
+  const { setAllowedSecondaryChainIds } = useWalletChainPolicy();
 
   const l1BlockNumber = executionBlockNumber ? BigInt(executionBlockNumber) : undefined;
   const l1TxHash = executorTxHash as Hex | undefined;
@@ -69,6 +72,15 @@ export function ProposalL2Execution({
   const detectedFromTx = !!message;
 
   const hasL2Leg = detectedFromActions || detectedFromTx;
+  const shouldAllowTaikoL2 = executed && shouldCheckL2 && !isL2Confirmed;
+
+  useEffect(() => {
+    setAllowedSecondaryChainIds(shouldAllowTaikoL2 ? [TAIKO_L2_CHAIN_ID] : []);
+
+    return () => {
+      setAllowedSecondaryChainIds([]);
+    };
+  }, [setAllowedSecondaryChainIds, shouldAllowTaikoL2]);
 
   // Don't render if no L2 leg detected (and not still checking)
   if (!shouldRenderL2ExecutionCard({ hasL2Leg, isExtracting, shouldCheckExecutedProposal: shouldCheckL2 })) {

packages/ui/src/context/WalletChainPolicy.tsx

@@ -0,0 +1,29 @@
+import React, { createContext, useContext, useMemo, useState } from "react";
+
+export interface WalletChainPolicyContextProps {
+  allowedSecondaryChainIds: number[];
+  setAllowedSecondaryChainIds: (chainIds: number[]) => void;
+}
+
+const WalletChainPolicyContext = createContext<WalletChainPolicyContextProps | undefined>(undefined);
+
+export const WalletChainPolicyProvider: React.FC<{ children: React.ReactNode }> = ({ children }) => {
+  const [allowedSecondaryChainIds, setAllowedSecondaryChainIds] = useState<number[]>([]);
+
+  const value = useMemo(
+    () => ({ allowedSecondaryChainIds, setAllowedSecondaryChainIds }),
+    [allowedSecondaryChainIds]
+  );
+
+  return <WalletChainPolicyContext.Provider value={value}>{children}</WalletChainPolicyContext.Provider>;
+};
+
+export const useWalletChainPolicy = () => {
+  const context = useContext(WalletChainPolicyContext);
+
+  if (!context) {
+    throw new Error("useWalletChainPolicy must be used inside the WalletChainPolicyProvider");
+  }
+
+  return context;
+};

packages/ui/src/context/index.tsx

@@ -11,6 +11,7 @@ import { PersistQueryClientProvider } from "@tanstack/react-query-persist-client
 import { UseDerivedWalletProvider } from "../hooks/useDerivedWallet";
 import { OdsModulesProvider } from "@aragon/ods";
 import { customModulesCopy, odsCoreProviderValues } from "@/components/ods-customizations";
+import { WalletChainPolicyProvider } from "./WalletChainPolicy";
 
 const queryClient = new QueryClient({
   defaultOptions: {
@@ -53,7 +54,9 @@ export function RootContextProvider({ children }: { children: ReactNode }) {
       >
         <PersistQueryClientProvider client={queryClient} persistOptions={{ persister }}>
           <AlertProvider>
-            <UseDerivedWalletProvider>{children}</UseDerivedWalletProvider>
+            <WalletChainPolicyProvider>
+              <UseDerivedWalletProvider>{children}</UseDerivedWalletProvider>
+            </WalletChainPolicyProvider>
           </AlertProvider>
         </PersistQueryClientProvider>
       </OdsModulesProvider>

packages/ui/src/tests/wallet-chain-policy.test.ts

@@ -0,0 +1,16 @@
+import { describe, expect, test } from "bun:test";
+import { shouldForcePrimaryChainSwitch } from "../utils/wallet-chain-policy";
+
+describe("shouldForcePrimaryChainSwitch", () => {
+  test("force-switches when connected to Taiko but the current flow has not allowed it", () => {
+    expect(shouldForcePrimaryChainSwitch(167000, 1, [])).toBe(true);
+  });
+
+  test("does not force-switch when the current flow explicitly allows Taiko", () => {
+    expect(shouldForcePrimaryChainSwitch(167000, 1, [167000])).toBe(false);
+  });
+
+  test("does not force-switch when already on the primary chain", () => {
+    expect(shouldForcePrimaryChainSwitch(1, 1, [])).toBe(false);
+  });
+});

packages/ui/src/utils/wallet-chain-policy.ts

@@ -0,0 +1,10 @@
+export function shouldForcePrimaryChainSwitch(
+  currentChainId: number | undefined,
+  primaryChainId: number,
+  allowedSecondaryChainIds: number[] = []
+) {
+  if (!currentChainId) return false;
+  if (currentChainId === primaryChainId) return false;
+
+  return !allowedSecondaryChainIds.includes(currentChainId);
+}