Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 18 additions & 4 deletions src/context/privacyBridge/usePrivacyBridgeWagmiSync.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand Down Expand Up @@ -45,6 +45,7 @@ export const usePrivacyBridgeWagmiSync = ({
setSessionAesKey,
setArePrivateBalancesHidden,
setPrivateTokens,
setPublicTokens,
checkSnapStatus,
clearSnapCache,
setMetamaskDetected,
Expand Down Expand Up @@ -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,
Expand All @@ -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,
]);
};
31 changes: 31 additions & 0 deletions tests/context/usePrivacyBridgeWagmiSync.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 () => {
Expand Down Expand Up @@ -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 () => {
Expand Down
Loading