From 970a51641cce99e11623bd8eaba8cb5604bda216 Mon Sep 17 00:00:00 2001
From: SimonK <simonk@coti.io>
Date: Thu, 30 Jul 2026 12:42:47 +0300
Subject: [PATCH] fix: reset public/private token lists on chain switch so Fuji
 metadata cannot stick on COTI.

---
 .../usePrivacyBridgeWagmiSync.ts              | 22 ++++++++++---
 .../context/usePrivacyBridgeWagmiSync.test.ts | 31 +++++++++++++++++++
 2 files changed, 49 insertions(+), 4 deletions(-)

diff --git a/src/context/privacyBridge/usePrivacyBridgeWagmiSync.ts b/src/context/privacyBridge/usePrivacyBridgeWagmiSync.ts
index f4a3971..ab6416b 100644
--- a/src/context/privacyBridge/usePrivacyBridgeWagmiSync.ts
+++ b/src/context/privacyBridge/usePrivacyBridgeWagmiSync.ts
@@ -5,7 +5,7 @@ import { isChainUpdatesMuted } from '../../lib/chainMute';
 import { logger } from '../../lib/logger';
 import { truncateAddress } from '../../lib/format';
 import { clearAesKeyValidatedForUnlock } from '../../crypto/aesKeyValidation';
-import { getInitialPrivateTokens } from '../../hooks/usePrivacyBridge';
+import { getInitialPrivateTokens, getInitialPublicTokens } from '../../hooks/usePrivacyBridge';
 import { reportPluginError, hasCotiErrorCode, CotiErrorCode } from '../../errors';
 import { isRateLimitedRpcError } from '../../lib/rpcProvider';
 import type { PrivacyBridgeAccountSync } from './usePrivacyBridgeAccountSync';
@@ -45,6 +45,7 @@ export const usePrivacyBridgeWagmiSync = ({
     setSessionAesKey,
     setArePrivateBalancesHidden,
     setPrivateTokens,
+    setPublicTokens,
     checkSnapStatus,
     clearSnapCache,
     setMetamaskDetected,
@@ -153,8 +154,11 @@ export const usePrivacyBridgeWagmiSync = ({
           prevWagmiChainIdRef.current = wagmiChainId;
           return;
         }
-        // When a session AES key exists, private balances are already correct.
-        // Re-fetch with the key to avoid resetting private tokens to zero.
+        // Swap both lists to the new chain together so Index pairing never
+        // mixes Fuji publics with COTI privates (or the reverse) mid-refresh.
+        setPublicTokens(getInitialPublicTokens(wagmiChainId));
+        setPrivateTokens(getInitialPrivateTokens(wagmiChainId));
+        // When a session AES key exists, re-fetch with the key so balances refill.
         if (core.sessionAesKey) {
           logger.log('[ChainChange] sessionAesKey present — refreshing with private balances', {
             from: prevWagmiChainIdRef.current,
@@ -176,5 +180,15 @@ export const usePrivacyBridgeWagmiSync = ({
       }
       prevWagmiChainIdRef.current = wagmiChainId;
     }
-  }, [wagmiConnected, wagmiAddress, walletAddress, isConnected, wagmiChainId, updateAccountState, core.sessionAesKey]);
+  }, [
+    wagmiConnected,
+    wagmiAddress,
+    walletAddress,
+    isConnected,
+    wagmiChainId,
+    updateAccountState,
+    core.sessionAesKey,
+    setPrivateTokens,
+    setPublicTokens,
+  ]);
 };
diff --git a/tests/context/usePrivacyBridgeWagmiSync.test.ts b/tests/context/usePrivacyBridgeWagmiSync.test.ts
index 9ef2f5b..b3d53db 100644
--- a/tests/context/usePrivacyBridgeWagmiSync.test.ts
+++ b/tests/context/usePrivacyBridgeWagmiSync.test.ts
@@ -197,6 +197,8 @@ describe('usePrivacyBridgeWagmiSync — chain-change guard with sessionAesKey',
         7082400,
       );
     });
+    expect(core.setPrivateTokens).toHaveBeenCalled();
+    expect(core.setPublicTokens).toHaveBeenCalled();
   });
 
   it('calls updateAccountState with fetchPrivate=false when sessionAesKey is null and chain changes (regression)', async () => {
@@ -225,6 +227,35 @@ describe('usePrivacyBridgeWagmiSync — chain-change guard with sessionAesKey',
         7082400,
       );
     });
+    expect(core.setPrivateTokens).toHaveBeenCalled();
+    expect(core.setPublicTokens).toHaveBeenCalled();
+  });
+
+  it('does not rewrite token lists when chain updates are muted', async () => {
+    h.isChainUpdatesMuted.mockReturnValue(true);
+
+    const core = makeCore({ sessionAesKey: 'c'.repeat(32), walletAddress: '0xabc123' });
+    const network = makeNetwork({ wagmiChainId: 11155111 });
+    const accountSync = makeAccountSync();
+
+    const { rerender } = renderHook(
+      (props) => usePrivacyBridgeWagmiSync(props),
+      { initialProps: { core, network, accountSync } },
+    );
+
+    vi.mocked(core.setPrivateTokens).mockClear();
+    vi.mocked(core.setPublicTokens).mockClear();
+
+    h.wagmiAccount.chainId = 7082400;
+    const updatedNetwork = makeNetwork({ wagmiChainId: 7082400 });
+
+    rerender({ core, network: updatedNetwork, accountSync });
+
+    await new Promise(r => setTimeout(r, 50));
+
+    // Muted onboarding must not rewrite token lists mid-flow
+    expect(core.setPrivateTokens).not.toHaveBeenCalled();
+    expect(core.setPublicTokens).not.toHaveBeenCalled();
   });
 
   it('does NOT call updateAccountState when chain updates are muted', async () => {