feat: add onConnect handler to <ConnectWallet />

package.json

@@ -34,7 +34,7 @@
     "react-dom": "^18"
   },
   "dependencies": {
-    "@rainbow-me/rainbowkit": "^2.1.3",
+    "@rainbow-me/rainbowkit": "^2.2.0",
     "@tanstack/react-query": "^5",
     "clsx": "^2.1.1",
     "graphql": "^14 || ^15 || ^16",

playground/nextjs-app-router/components/demo/Wallet.tsx

@@ -15,18 +15,37 @@ import {
   WalletDropdownFundLink,
   WalletDropdownLink,
 } from '@coinbase/onchainkit/wallet';
-import { useAccount } from 'wagmi';
+import { useAccount, useSignMessage } from 'wagmi';
+import { createSiweMessage } from 'viem/siwe';
+
+import '@rainbow-me/rainbowkit/styles.css';
+
+const message = createSiweMessage({
+  address: '0xae9eCa1Fa2F786E16D35F5B8C5df3Ac484c490FF',
+  chainId: 1,
+  domain: 'example.com',
+  nonce: 'foobarbaz',
+  uri: 'https://example.com/path',
+  version: '1',
+});
 
 function WalletComponent() {
   const { address } = useAccount();
+  const { signMessage } = useSignMessage();
 
   return (
     <div className="flex justify-end">
       <Wallet>
-        <ConnectWallet text="Connect Wallet">
+        <ConnectWallet
+          text="Connect Wallet"
+          onInitialConnect={() => {
+            signMessage({ message });
+          }}
+        >
           <Avatar address={address} className="h-6 w-6" />
           <Name />
         </ConnectWallet>
+
         <WalletDropdown>
           <Identity className="px-4 pt-3 pb-2" hasCopyAddressOnClick={true}>
             <Avatar />

playground/nextjs-app-router/onchainkit/package.json

@@ -34,7 +34,7 @@
     "react-dom": "^18"
   },
   "dependencies": {
-    "@rainbow-me/rainbowkit": "^2.1.3",
+    "@rainbow-me/rainbowkit": "^2.2.0",
     "@tanstack/react-query": "^5",
     "clsx": "^2.1.1",
     "graphql": "^14 || ^15 || ^16",

playground/nextjs-app-router/package.json

@@ -37,7 +37,7 @@
     "react-dom": "^18",
     "tailwind-merge": "^2.4.0",
     "tailwindcss-animate": "^1.0.7",
-    "viem": "^2.17.4",
+    "viem": "^2.21.35",
     "wagmi": "^2.11.0"
   },
   "devDependencies": {

src/wallet/components/ConnectWallet.test.tsx

@@ -88,7 +88,7 @@ describe('ConnectWallet', () => {
     expect(connectedText).toBeInTheDocument();
   });
 
-  it('should calls connect function when connect button is clicked', () => {
+  it('should call connect function when connect button is clicked', () => {
     const connectMock = vi.fn();
     vi.mocked(useConnect).mockReturnValue({
       connectors: [{ id: 'mockConnector' }],
@@ -98,9 +98,14 @@ describe('ConnectWallet', () => {
     render(<ConnectWallet text="Connect Wallet" />);
     const button = screen.getByTestId('ockConnectButton');
     fireEvent.click(button);
-    expect(connectMock).toHaveBeenCalledWith({
-      connector: { id: 'mockConnector' },
-    });
+    expect(connectMock).toHaveBeenCalledWith(
+      {
+        connector: { id: 'mockConnector' },
+      },
+      {
+        onSuccess: expect.any(Function),
+      },
+    );
   });
 
   it('should toggle wallet modal on button click when connected', () => {
@@ -162,6 +167,71 @@ describe('ConnectWallet', () => {
     expect(screen.queryByText('Not Render')).not.toBeInTheDocument();
   });
 
+  it('should call onInitialConnect callback when connect button is clicked', async () => {
+    const mockUseAccount = vi.mocked(useAccount);
+    const connectMock = vi.fn();
+    const onInitialConnectMock = vi.fn();
+
+    // Initial state: disconnected
+    mockUseAccount.mockReturnValue({
+      address: undefined,
+      status: 'disconnected',
+    });
+
+    vi.mocked(useConnect).mockReturnValue({
+      connectors: [{ id: 'mockConnector' }],
+      connect: connectMock,
+      status: 'idle',
+    });
+
+    render(
+      <ConnectWallet
+        text="Connect Wallet"
+        onInitialConnect={onInitialConnectMock}
+      />,
+    );
+
+    const button = screen.getByTestId('ockConnectButton');
+    fireEvent.click(button);
+
+    // Simulate successful connection
+    connectMock.mock.calls[0][1].onSuccess();
+
+    // Update account status to connected
+    mockUseAccount.mockReturnValue({
+      address: '0x123',
+      status: 'connected',
+    });
+
+    // Force a re-render to trigger the useEffect
+    render(
+      <ConnectWallet
+        text="Connect Wallet"
+        onInitialConnect={onInitialConnectMock}
+      />,
+    );
+
+    expect(onInitialConnectMock).toHaveBeenCalledTimes(1);
+  });
+
+  it('should not call onConnect callback when component is first mounted', () => {
+    const mockUseAccount = vi.mocked(useAccount);
+    mockUseAccount.mockReturnValue({
+      address: '0x123',
+      status: 'connected',
+    });
+
+    const onInitialConnectMock = vi.fn();
+    render(
+      <ConnectWallet
+        text="Connect Wallet"
+        onInitialConnect={onInitialConnectMock}
+      />,
+    );
+
+    expect(onInitialConnectMock).toHaveBeenCalledTimes(0);
+  });
+
   describe('withWalletAggregator', () => {
     beforeEach(() => {
       vi.mocked(useAccount).mockReturnValue({
@@ -175,7 +245,7 @@ describe('ConnectWallet', () => {
       });
     });
 
-    it('should render ConnectButtonRainboKit when withWalletAggregator is true', () => {
+    it('should render ConnectButtonRainbowKit when withWalletAggregator is true', () => {
       render(
         <ConnectWallet text="Connect Wallet" withWalletAggregator={true} />,
       );
@@ -198,12 +268,17 @@ describe('ConnectWallet', () => {
       );
       const connectButton = screen.getByTestId('ockConnectButton');
       fireEvent.click(connectButton);
-      expect(connectMock).toHaveBeenCalledWith({
-        connector: { id: 'mockConnector' },
-      });
+      expect(connectMock).toHaveBeenCalledWith(
+        {
+          connector: { id: 'mockConnector' },
+        },
+        {
+          onSuccess: expect.any(Function),
+        },
+      );
     });
 
-    it('should calls openConnectModal function when connect button is clicked', () => {
+    it('should call openConnectModal function when connect button is clicked', () => {
       vi.mocked(useWalletContext).mockReturnValue({
         isOpen: false,
         setIsOpen: vi.fn(),
@@ -215,5 +290,32 @@ describe('ConnectWallet', () => {
       fireEvent.click(button);
       expect(openConnectModalMock).toHaveBeenCalled();
     });
+
+    it('should call onConnect callback when connect button is clicked', () => {
+      const mockUseAccount = vi.mocked(useAccount);
+      mockUseAccount.mockReturnValue({
+        address: undefined,
+        status: 'disconnected',
+      });
+
+      const onInitialConnectMock = vi.fn();
+      render(
+        <ConnectWallet
+          text="Connect Wallet"
+          onInitialConnect={onInitialConnectMock}
+          withWalletAggregator={true}
+        />,
+      );
+      const button = screen.getByTestId('ockConnectButton');
+
+      mockUseAccount.mockReturnValue({
+        address: '0x123',
+        status: 'connected',
+      });
+
+      fireEvent.click(button);
+
+      expect(onInitialConnectMock).toHaveBeenCalledTimes(1);
+    });
   });
 });

src/wallet/components/ConnectWallet.tsx

@@ -1,6 +1,7 @@
 import { ConnectButton as ConnectButtonRainbowKit } from '@rainbow-me/rainbowkit';
 import { Children, isValidElement, useCallback, useMemo } from 'react';
 import type { ReactNode } from 'react';
+import React from 'react';
 import { useAccount, useConnect } from 'wagmi';
 import { IdentityProvider } from '../../identity/components/IdentityProvider';
 import { Spinner } from '../../internal/components/Spinner';
@@ -24,12 +25,16 @@ export function ConnectWallet({
   // but for now we will keep it for backward compatibility.
   text = 'Connect Wallet',
   withWalletAggregator = false,
+  onInitialConnect,
 }: ConnectWalletReact) {
   // Core Hooks
   const { isOpen, setIsOpen } = useWalletContext();
   const { address: accountAddress, status } = useAccount();
   const { connectors, connect, status: connectStatus } = useConnect();
 
+  // State
+  const [hasClickedConnect, setHasClickedConnect] = React.useState(false);
+
   // Get connectWalletText from children when present,
   // this is used to customize the connect wallet button text
   const { connectWalletText } = useMemo(() => {
@@ -58,6 +63,14 @@ export function ConnectWallet({
     setIsOpen(!isOpen);
   }, [isOpen, setIsOpen]);
 
+  // Effects
+  React.useEffect(() => {
+    if (hasClickedConnect && status === 'connected' && onInitialConnect) {
+      onInitialConnect();
+      setHasClickedConnect(false);
+    }
+  }, [status, hasClickedConnect, onInitialConnect]);
+
   if (status === 'disconnected') {
     if (withWalletAggregator) {
       return (
@@ -67,7 +80,10 @@ export function ConnectWallet({
               <ConnectButton
                 className={className}
                 connectWalletText={connectWalletText}
-                onClick={() => openConnectModal()}
+                onClick={() => {
+                  openConnectModal();
+                  setHasClickedConnect(true);
+                }}
                 text={text}
               />
             </div>
@@ -80,7 +96,16 @@ export function ConnectWallet({
         <ConnectButton
           className={className}
           connectWalletText={connectWalletText}
-          onClick={() => connect({ connector })}
+          onClick={() => {
+            connect(
+              { connector },
+              {
+                onSuccess: () => {
+                  onInitialConnect?.();
+                },
+              },
+            );
+          }}
           text={text}
         />
       </div>

src/wallet/types.ts

@@ -8,6 +8,7 @@ export type ConnectButtonReact = {
   className?: string; // Optional className override for button element
   connectWalletText: ReactNode | null; // Optional text override for button
   onClick: () => void; // Function to call when the button is clicked
+  onConnect?: () => void; // Optional callback function to execute when the wallet is connected
   text: string; // Optional text override for button
 };
 
@@ -20,6 +21,11 @@ export type ConnectWalletReact = {
   /** @deprecated Prefer `ConnectWalletText component` */
   text?: string; // Optional text override for button
   withWalletAggregator?: boolean; // Optional flag to enable the wallet aggregator like RainbowKit
+  /**
+   * Optional callback function to execute when the wallet is connected.
+   * This can be used to perform additional actions after the wallet is connected.
+   */
+  onInitialConnect?: () => void;
 };
 
 /**