fix: better UX for off-chain message signing

src/components/tx-flow/flows/SignMessage/SignMessage.test.tsx

@@ -16,6 +16,7 @@ import { render, act, fireEvent, waitFor } from '@/tests/test-utils'
 import type { ConnectedWallet } from '@/hooks/wallets/useOnboard'
 import type { EIP1193Provider, WalletState, AppState, OnboardAPI } from '@web3-onboard/core'
 import { generateSafeMessageHash } from '@/utils/safe-messages'
+import { getSafeMessage } from '@safe-global/safe-gateway-typescript-sdk'
 
 jest.mock('@safe-global/safe-gateway-typescript-sdk', () => ({
   ...jest.requireActual('@safe-global/safe-gateway-typescript-sdk'),
@@ -68,6 +69,14 @@ const mockOnboard = {
 } as unknown as OnboardAPI
 
 describe('SignMessage', () => {
+  beforeAll(() => {
+    jest.useFakeTimers()
+  })
+
+  afterAll(() => {
+    jest.useRealTimers()
+  })
+
   beforeEach(() => {
     jest.clearAllMocks()
 
@@ -196,7 +205,7 @@ describe('SignMessage', () => {
 
     jest.spyOn(useAsyncHook, 'default').mockReturnValue([undefined, new Error('SafeMessage not found'), false])
 
-    const { getByText } = render(
+    const { getByText, baseElement } = render(
       <SignMessage
         logoUri="www.fake.com/test.png"
         name="Test App"
@@ -206,7 +215,23 @@ describe('SignMessage', () => {
       />,
     )
 
-    const proposalSpy = jest.spyOn(sender, 'dispatchSafeMsgProposal')
+    const proposalSpy = jest.spyOn(sender, 'dispatchSafeMsgProposal').mockImplementation(() => Promise.resolve())
+    const mockMessageHash = '0x456'
+    const msg = {
+      type: SafeMessageListItemType.MESSAGE,
+      messageHash: mockMessageHash,
+      confirmations: [
+        {
+          owner: {
+            value: hexZeroPad('0x2', 20),
+          },
+        },
+      ],
+      confirmationsRequired: 2,
+      confirmationsSubmitted: 1,
+    } as unknown as SafeMessage
+
+    ;(getSafeMessage as jest.Mock).mockResolvedValue(msg)
 
     const button = getByText('Sign')
 
@@ -228,6 +253,11 @@ describe('SignMessage', () => {
         safeAppId: 25,
       }),
     )
+
+    // Immediately refetches message and displays confirmation
+    expect(baseElement).toHaveTextContent('0x0000...0002')
+    expect(baseElement).toHaveTextContent('1 of 2')
+    expect(baseElement).toHaveTextContent('Confirmation #2')
   })
 
   it('confirms the message if already proposed', async () => {
@@ -275,11 +305,31 @@ describe('SignMessage', () => {
       Promise.resolve()
     })
 
-    const confirmationSpy = jest.spyOn(sender, 'dispatchSafeMsgConfirmation')
+    const confirmationSpy = jest
+      .spyOn(sender, 'dispatchSafeMsgConfirmation')
+      .mockImplementation(() => Promise.resolve())
 
     const button = getByText('Sign')
 
     expect(button).toBeEnabled()
+    ;(getSafeMessage as jest.Mock).mockResolvedValue({
+      ...msg,
+      confirmations: [
+        {
+          owner: {
+            value: hexZeroPad('0x2', 20),
+          },
+        },
+        {
+          owner: {
+            value: hexZeroPad('0x3', 20),
+          },
+        },
+      ],
+      confirmationsRequired: 2,
+      confirmationsSubmitted: 2,
+      preparedSignature: '0x789',
+    })
 
     await act(() => {
       fireEvent.click(button)
@@ -298,6 +348,12 @@ describe('SignMessage', () => {
         message: 'Hello world!',
       }),
     )
+
+    await act(() => {
+      jest.advanceTimersByTime(1000)
+    })
+
+    expect(getByText('Message successfully signed')).toBeInTheDocument()
   })
 
   it('displays an error if no wallet is connected', () => {
@@ -448,6 +504,10 @@ describe('SignMessage', () => {
       fireEvent.click(button)
     })
 
+    await act(() => {
+      jest.advanceTimersByTime(1000)
+    })
+
     await waitFor(() => {
       expect(proposalSpy).toHaveBeenCalled()
       expect(getByText('Error confirming the message. Please try again.')).toBeInTheDocument()
@@ -488,6 +548,11 @@ describe('SignMessage', () => {
       fireEvent.click(button)
     })
 
+    await act(() => {
+      jest.advanceTimersByTime(1000)
+    })
+
+    // We automatically call the dispatch after 1 second
     expect(confirmationSpy).toHaveBeenCalled()
 
     await waitFor(() => {

src/components/tx-flow/flows/SignMessage/SignMessage.tsx

@@ -1,6 +1,18 @@
-import { Grid, Button, Box, Typography, SvgIcon, CardContent, CardActions } from '@mui/material'
+import {
+  Grid,
+  Button,
+  Box,
+  Typography,
+  SvgIcon,
+  CardContent,
+  CardActions,
+  Accordion,
+  AccordionSummary,
+  AccordionDetails,
+} from '@mui/material'
 import { useTheme } from '@mui/material/styles'
-import { useContext } from 'react'
+import ExpandMoreIcon from '@mui/icons-material/ExpandMore'
+import { useContext, useEffect, useState } from 'react'
 import { SafeMessageListItemType, SafeMessageStatus } from '@safe-global/safe-gateway-typescript-sdk'
 import type { ReactElement } from 'react'
 import type { SafeMessage } from '@safe-global/safe-gateway-typescript-sdk'
@@ -25,6 +37,7 @@ import useHighlightHiddenTab from '@/hooks/useHighlightHiddenTab'
 import InfoBox from '@/components/safe-messages/InfoBox'
 import { DecodedMsg } from '@/components/safe-messages/DecodedMsg'
 import TxCard from '@/components/tx-flow/common/TxCard'
+import { dispatchPreparedSignature } from '@/services/safe-messages/safeMsgNotifications'
 
 const createSkeletonMessage = (confirmationsRequired: number): SafeMessage => {
   return {
@@ -64,9 +77,11 @@ const DialogHeader = ({ threshold }: { threshold: number }) => (
     <Typography variant="h4" textAlign="center" gutterBottom>
       Confirm message
     </Typography>
-    <Typography variant="body1" textAlign="center" mb={2}>
-      To sign this message, you need to collect <b>{threshold} owner signatures</b> of your Safe Account.
-    </Typography>
+    {threshold > 1 && (
+      <Typography variant="body1" textAlign="center" mb={2}>
+        To sign this message, collect signatures from <b>{threshold} owners</b> of your Safe Account.
+      </Typography>
+    )}
   </>
 )
 
@@ -116,6 +131,23 @@ const AlreadySignedByOwnerMessage = ({ hasSigned }: { hasSigned: boolean }) => {
   )
 }
 
+const SuccessCard = ({ safeMessage, onContinue }: { safeMessage: SafeMessage; onContinue: () => void }) => {
+  return (
+    <TxCard>
+      <Typography variant="h4" textAlign="center" gutterBottom>
+        Message successfully signed
+      </Typography>
+
+      <MsgSigners msg={safeMessage} showOnlyConfirmations showMissingSignatures />
+      <CardActions>
+        <Button variant="contained" color="primary" onClick={onContinue} disabled={!safeMessage.preparedSignature}>
+          Continue
+        </Button>
+      </CardActions>
+    </TxCard>
+  )
+}
+
 type BaseProps = Pick<SafeMessage, 'logoUri' | 'name' | 'message'>
 
 // Custom Safe Apps do not have a `safeAppId`
@@ -140,24 +172,47 @@ const SignMessage = ({ message, safeAppId, requestId }: ProposeProps | ConfirmPr
 
   const { decodedMessage, safeMessageMessage, safeMessageHash } = useDecodedSafeMessage(message, safe)
   const ongoingMessage = useSafeMessage(safeMessageHash)
+  const [safeMessage, setSafeMessage] = useState(ongoingMessage)
+
+  // Sync ongoing msg
+  useEffect(() => {
+    setSafeMessage(ongoingMessage)
+  }, [ongoingMessage])
+
   useHighlightHiddenTab()
 
   const decodedMessageAsString =
     typeof decodedMessage === 'string' ? decodedMessage : JSON.stringify(decodedMessage, null, 2)
 
-  const hasSigned = !!ongoingMessage?.confirmations.some(({ owner }) => owner.value === wallet?.address)
+  const hasSigned = !!safeMessage?.confirmations.some(({ owner }) => owner.value === wallet?.address)
+
+  const isFullySigned = !!safeMessage?.preparedSignature
 
   const isDisabled = !isOwner || hasSigned
 
   const { onSign, submitError } = useSyncSafeMessageSigner(
-    ongoingMessage,
+    safeMessage,
     decodedMessage,
     safeMessageHash,
     requestId,
     safeAppId,
     () => setTxFlow(undefined),
   )
 
+  const handleSign = async () => {
+    const updatedMessage = await onSign()
+    if (updatedMessage) {
+      setSafeMessage(updatedMessage)
+    }
+  }
+
+  const onContinue = async () => {
+    if (!safeMessage) {
+      return
+    }
+    await dispatchPreparedSignature(safeMessage, safeMessageHash, () => setTxFlow(undefined), requestId)
+  }
+
   return (
     <>
       <TxCard>
@@ -169,42 +224,52 @@ const SignMessage = ({ message, safeAppId, requestId }: ProposeProps | ConfirmPr
           </Typography>
           <DecodedMsg message={decodedMessage} isInModal />
 
-          <MessageHashField label="SafeMessage" hashValue={safeMessageMessage} />
-          <MessageHashField label="SafeMessage hash" hashValue={safeMessageHash} />
+          <Accordion sx={{ mt: 2, '&.Mui-expanded': { mt: 2 } }}>
+            <AccordionSummary expandIcon={<ExpandMoreIcon />}>SafeMessage details</AccordionSummary>
+            <AccordionDetails>
+              <MessageHashField label="SafeMessage" hashValue={safeMessageMessage} />
+              <MessageHashField label="SafeMessage hash" hashValue={safeMessageHash} />
+            </AccordionDetails>
+          </Accordion>
         </CardContent>
       </TxCard>
 
-      <TxCard>
-        <AlreadySignedByOwnerMessage hasSigned={hasSigned} />
-
-        <InfoBox
-          title="Collect all the confirmations"
-          message={
-            requestId
-              ? 'Please keep this modal open until all signers confirm this message. Closing the modal will abort the signing request.'
-              : 'The signature will be submitted to the Safe App when the message is fully signed.'
-          }
-        >
-          <MsgSigners
-            msg={ongoingMessage || createSkeletonMessage(safe.threshold)}
-            showOnlyConfirmations
-            showMissingSignatures
-            backgroundColor={palette.info.background}
-          />
-        </InfoBox>
-
-        <WrongChainWarning />
-
-        <MessageDialogError isOwner={isOwner} submitError={submitError} />
-      </TxCard>
+      {isFullySigned ? (
+        <SuccessCard onContinue={onContinue} safeMessage={safeMessage} />
+      ) : (
+        <>
+          <TxCard>
+            <AlreadySignedByOwnerMessage hasSigned={hasSigned} />
 
-      <TxCard>
-        <CardActions>
-          <Button variant="contained" color="primary" onClick={onSign} disabled={isDisabled}>
-            Sign
-          </Button>
-        </CardActions>
-      </TxCard>
+            <InfoBox
+              title="Collect all the confirmations"
+              message={
+                requestId
+                  ? 'Please keep this modal open until all signers confirm this message. Closing the modal will abort the signing request.'
+                  : 'The signature will be submitted to the Safe App when the message is fully signed.'
+              }
+            >
+              <MsgSigners
+                msg={safeMessage || createSkeletonMessage(safe.threshold)}
+                showOnlyConfirmations
+                showMissingSignatures
+                backgroundColor={palette.info.background}
+              />
+            </InfoBox>
+
+            <WrongChainWarning />
+
+            <MessageDialogError isOwner={isOwner} submitError={submitError} />
+          </TxCard>
+          <TxCard>
+            <CardActions>
+              <Button variant="contained" color="primary" onClick={handleSign} disabled={isDisabled}>
+                Sign
+              </Button>
+            </CardActions>
+          </TxCard>
+        </>
+      )}
     </>
   )
 }