feat: optimistic disconnect

.changeset/wild-hounds-talk.md

@@ -0,0 +1,5 @@
+---
+"@rainbow-me/rainbowkit": patch
+---
+
+Introduced optimistic disconnect behavior by monitoring `isDisconnecting` in `RainbowKitWagmiStateProvider`. If `isDisconnecting` is `true`, the `useConnectionStatus` hook will mark its state as `disconnected`.

packages/rainbowkit/src/components/AccountModal/AccountModal.tsx

@@ -5,6 +5,7 @@ import { useMainnetEnsName } from '../../hooks/useMainnetEnsName';
 import { Dialog } from '../Dialog/Dialog';
 import { DialogContent } from '../Dialog/DialogContent';
 import { ProfileDetails } from '../ProfileDetails/ProfileDetails';
+import { useRainbowKitWagmiState } from '../RainbowKitProvider/RainbowKitWagmiStateProvider';
 
 export interface AccountModalProps {
   open: boolean;
@@ -15,7 +16,14 @@ export function AccountModal({ onClose, open }: AccountModalProps) {
   const { address } = useAccount();
   const ensName = useMainnetEnsName(address);
   const ensAvatar = useMainnetEnsAvatar(ensName);
-  const { disconnect } = useDisconnect();
+  const { setIsDisconnecting } = useRainbowKitWagmiState();
+
+  const { disconnect } = useDisconnect({
+    mutation: {
+      onMutate: () => setIsDisconnecting(true),
+      onSettled: () => setIsDisconnecting(false),
+    },
+  });
 
   if (!address) {
     return null;
@@ -33,7 +41,10 @@ export function AccountModal({ onClose, open }: AccountModalProps) {
               ensAvatar={ensAvatar}
               ensName={ensName}
               onClose={onClose}
-              onDisconnect={disconnect}
+              onDisconnect={() => {
+                onClose();
+                disconnect();
+              }}
             />
           </DialogContent>
         </Dialog>

packages/rainbowkit/src/components/ChainModal/ChainModal.tsx

@@ -10,6 +10,7 @@ import { DisconnectSqIcon } from '../Icons/DisconnectSq';
 import { MenuButton } from '../MenuButton/MenuButton';
 import { I18nContext } from '../RainbowKitProvider/I18nContext';
 import { useRainbowKitChains } from '../RainbowKitProvider/RainbowKitChainContext';
+import { useRainbowKitWagmiState } from '../RainbowKitProvider/RainbowKitWagmiStateProvider';
 import { Text } from '../Text/Text';
 import Chain from './Chain';
 import {
@@ -45,7 +46,14 @@ export function ChainModal({ onClose, open }: ChainModalProps) {
 
   const { i18n } = useContext(I18nContext);
 
-  const { disconnect } = useDisconnect();
+  const { disconnect } = useDisconnect({
+    mutation: {
+      onMutate: () => setIsDisconnecting(true),
+      onSettled: () => setIsDisconnecting(false),
+    },
+  });
+
+  const { setIsDisconnecting } = useRainbowKitWagmiState();
   const titleId = 'rk_chain_modal_title';
   const mobile = isMobile();
   const isCurrentChainSupported = chains.some((chain) => chain.id === chainId);
@@ -116,7 +124,10 @@ export function ChainModal({ onClose, open }: ChainModalProps) {
               <>
                 <Box background="generalBorderDim" height="1" marginX="8" />
                 <MenuButton
-                  onClick={() => disconnect()}
+                  onClick={() => {
+                    onClose();
+                    disconnect();
+                  }}
                   testId="chain-option-disconnect"
                 >
                   <Box

packages/rainbowkit/src/components/ConnectModal/ConnectModal.tsx

@@ -4,6 +4,7 @@ import { useConnectionStatus } from '../../hooks/useConnectionStatus';
 import ConnectOptions from '../ConnectOptions/ConnectOptions';
 import { Dialog } from '../Dialog/Dialog';
 import { DialogContent } from '../Dialog/DialogContent';
+import { useRainbowKitWagmiState } from '../RainbowKitProvider/RainbowKitWagmiStateProvider';
 import { SignIn } from '../SignIn/SignIn';
 
 export interface ConnectModalProps {
@@ -14,8 +15,15 @@ export interface ConnectModalProps {
 export function ConnectModal({ onClose, open }: ConnectModalProps) {
   const titleId = 'rk_connect_title';
   const connectionStatus = useConnectionStatus();
+  const { setIsDisconnecting } = useRainbowKitWagmiState();
+
+  const { disconnect } = useDisconnect({
+    mutation: {
+      onMutate: () => setIsDisconnecting(true),
+      onSettled: () => setIsDisconnecting(false),
+    },
+  });
 
-  const { disconnect } = useDisconnect();
   const { isConnecting } = useAccount();
 
   // when a user cancels or dismisses the SignIn modal for SIWE, disconnect and call onClose

packages/rainbowkit/src/components/RainbowKitProvider/ModalContext.tsx

@@ -13,6 +13,7 @@ import { AccountModal } from '../AccountModal/AccountModal';
 import { ChainModal } from '../ChainModal/ChainModal';
 import { ConnectModal } from '../ConnectModal/ConnectModal';
 import { useAuthenticationStatus } from './AuthenticationContext';
+import { useRainbowKitWagmiState } from './RainbowKitWagmiStateProvider';
 
 function useModalStateValue() {
   const [isModalOpen, setModalOpen] = useState(false);
@@ -71,6 +72,8 @@ export function ModalProvider({ children }: ModalProviderProps) {
 
   const connectionStatus = useConnectionStatus();
 
+  const { isDisconnecting } = useRainbowKitWagmiState();
+
   const { chainId } = useAccount();
   const { chains } = useConfig();
 
@@ -119,13 +122,18 @@ export function ModalProvider({ children }: ModalProviderProps) {
           openChainModal:
             connectionStatus === 'connected' ? openChainModal : undefined,
           openConnectModal:
-            connectionStatus === 'disconnected' ||
-            connectionStatus === 'unauthenticated'
+            // Prevent opening the connect modal during disconnecting mode. Even though `connectionStatus`
+            // may mark it as 'disconnected', we don't want the user to open the modal as the wagmi state
+            // still considers it 'connected'
+            !isDisconnecting &&
+            (connectionStatus === 'disconnected' ||
+              connectionStatus === 'unauthenticated')
               ? openConnectModal
               : undefined,
           setIsWalletConnectModalOpen,
         }),
         [
+          isDisconnecting,
           connectionStatus,
           accountModalOpen,
           chainModalOpen,

packages/rainbowkit/src/components/RainbowKitProvider/RainbowKitProvider.tsx

@@ -17,6 +17,7 @@ import {
   ModalSizes,
 } from './ModalSizeContext';
 import { RainbowKitChainProvider } from './RainbowKitChainContext';
+import { RainbowKitWagmiStateProvider } from './RainbowKitWagmiStateProvider';
 import { ShowBalanceProvider } from './ShowBalanceContext';
 import { ShowRecentTransactionsContext } from './ShowRecentTransactionsContext';
 import { WalletButtonProvider } from './WalletButtonContext';
@@ -102,61 +103,63 @@ export function RainbowKitProvider({
   const avatarContext = avatar ?? defaultAvatar;
 
   return (
-    <RainbowKitChainProvider initialChain={initialChain}>
-      <WalletButtonProvider>
-        <I18nProvider locale={locale}>
-          <CoolModeContext.Provider value={coolMode}>
-            <ModalSizeProvider modalSize={modalSize}>
-              <ShowRecentTransactionsContext.Provider
-                value={showRecentTransactions}
-              >
-                <TransactionStoreProvider>
-                  <AvatarContext.Provider value={avatarContext}>
-                    <AppContext.Provider value={appContext}>
-                      <ThemeIdContext.Provider value={id}>
-                        <ShowBalanceProvider>
-                          <ModalProvider>
-                            {theme ? (
-                              <div {...createThemeRootProps(id)}>
-                                <style
-                                  // biome-ignore lint/security/noDangerouslySetInnerHtml: TODO
-                                  dangerouslySetInnerHTML={{
-                                    // Selectors are sanitized to only contain alphanumeric
-                                    // and underscore characters. Theme values generated by
-                                    // cssStringFromTheme are sanitized, removing
-                                    // characters that terminate values / HTML tags.
-                                    __html: [
-                                      `${selector}{${cssStringFromTheme(
-                                        'lightMode' in theme
-                                          ? theme.lightMode
-                                          : theme,
-                                      )}}`,
-
-                                      'darkMode' in theme
-                                        ? `@media(prefers-color-scheme:dark){${selector}{${cssStringFromTheme(
-                                            theme.darkMode,
-                                            { extends: theme.lightMode },
-                                          )}}}`
-                                        : null,
-                                    ].join(''),
-                                  }}
-                                />
-                                {children}
-                              </div>
-                            ) : (
-                              children
-                            )}
-                          </ModalProvider>
-                        </ShowBalanceProvider>
-                      </ThemeIdContext.Provider>
-                    </AppContext.Provider>
-                  </AvatarContext.Provider>
-                </TransactionStoreProvider>
-              </ShowRecentTransactionsContext.Provider>
-            </ModalSizeProvider>
-          </CoolModeContext.Provider>
-        </I18nProvider>
-      </WalletButtonProvider>
-    </RainbowKitChainProvider>
+    <RainbowKitWagmiStateProvider>
+      <RainbowKitChainProvider initialChain={initialChain}>
+        <WalletButtonProvider>
+          <I18nProvider locale={locale}>
+            <CoolModeContext.Provider value={coolMode}>
+              <ModalSizeProvider modalSize={modalSize}>
+                <ShowRecentTransactionsContext.Provider
+                  value={showRecentTransactions}
+                >
+                  <TransactionStoreProvider>
+                    <AvatarContext.Provider value={avatarContext}>
+                      <AppContext.Provider value={appContext}>
+                        <ThemeIdContext.Provider value={id}>
+                          <ShowBalanceProvider>
+                            <ModalProvider>
+                              {theme ? (
+                                <div {...createThemeRootProps(id)}>
+                                  <style
+                                    // biome-ignore lint/security/noDangerouslySetInnerHtml: TODO
+                                    dangerouslySetInnerHTML={{
+                                      // Selectors are sanitized to only contain alphanumeric
+                                      // and underscore characters. Theme values generated by
+                                      // cssStringFromTheme are sanitized, removing
+                                      // characters that terminate values / HTML tags.
+                                      __html: [
+                                        `${selector}{${cssStringFromTheme(
+                                          'lightMode' in theme
+                                            ? theme.lightMode
+                                            : theme,
+                                        )}}`,
+
+                                        'darkMode' in theme
+                                          ? `@media(prefers-color-scheme:dark){${selector}{${cssStringFromTheme(
+                                              theme.darkMode,
+                                              { extends: theme.lightMode },
+                                            )}}}`
+                                          : null,
+                                      ].join(''),
+                                    }}
+                                  />
+                                  {children}
+                                </div>
+                              ) : (
+                                children
+                              )}
+                            </ModalProvider>
+                          </ShowBalanceProvider>
+                        </ThemeIdContext.Provider>
+                      </AppContext.Provider>
+                    </AvatarContext.Provider>
+                  </TransactionStoreProvider>
+                </ShowRecentTransactionsContext.Provider>
+              </ModalSizeProvider>
+            </CoolModeContext.Provider>
+          </I18nProvider>
+        </WalletButtonProvider>
+      </RainbowKitChainProvider>
+    </RainbowKitWagmiStateProvider>
   );
 }

packages/rainbowkit/src/components/RainbowKitProvider/RainbowKitWagmiStateProvider.tsx

@@ -0,0 +1,47 @@
+import React, {
+  ReactNode,
+  createContext,
+  useContext,
+  useMemo,
+  useState,
+} from 'react';
+
+interface RainbowKitWagmiStateContextValue {
+  isDisconnecting: boolean;
+  setIsDisconnecting: (isDisconnecting: boolean) => void;
+}
+
+const RainbowKitWagmiStateContext =
+  createContext<RainbowKitWagmiStateContextValue>({
+    isDisconnecting: false,
+    setIsDisconnecting: () => {},
+  });
+
+interface RainbowKitWagmiStateProviderProps {
+  children: ReactNode;
+}
+
+export function RainbowKitWagmiStateProvider({
+  children,
+}: RainbowKitWagmiStateProviderProps) {
+  // The 'status' state from the `useDisconnect` hook in wagmi can't be used for 'pending' logic,
+  // as it doesn't share its state globally. It only shares its state within the same component (like useState).
+  const [isDisconnecting, setIsDisconnecting] = useState(false);
+
+  return (
+    <RainbowKitWagmiStateContext.Provider
+      value={useMemo(
+        () => ({
+          isDisconnecting,
+          setIsDisconnecting,
+        }),
+        [isDisconnecting],
+      )}
+    >
+      {children}
+    </RainbowKitWagmiStateContext.Provider>
+  );
+}
+
+export const useRainbowKitWagmiState = () =>
+  useContext(RainbowKitWagmiStateContext);

packages/rainbowkit/src/hooks/useConnectionStatus.ts

@@ -1,5 +1,6 @@
 import { useAccount } from 'wagmi';
 import { useAuthenticationStatus } from '../components/RainbowKitProvider/AuthenticationContext';
+import { useRainbowKitWagmiState } from '../components/RainbowKitProvider/RainbowKitWagmiStateProvider';
 
 export type ConnectionStatus =
   | 'disconnected'
@@ -10,8 +11,9 @@ export type ConnectionStatus =
 export function useConnectionStatus(): ConnectionStatus {
   const authenticationStatus = useAuthenticationStatus();
   const { isConnected } = useAccount();
+  const { isDisconnecting } = useRainbowKitWagmiState();
 
-  if (!isConnected) {
+  if (!isConnected || isDisconnecting) {
     return 'disconnected';
   }