fix(pilot-app, extension): better optimistic updates and connection handling

deployables/app/app/components/wallet/ConnectWalletFallback.tsx

@@ -0,0 +1,9 @@
+import { Labeled, PrimaryButton } from '@zodiac/ui'
+
+export const ConnectWalletFallback = () => (
+  <Labeled label="Pilot Account">
+    <PrimaryButton fluid disabled>
+      Loading wallet support...
+    </PrimaryButton>
+  </Labeled>
+)

deployables/app/app/components/wallet/WalletProvider.tsx

@@ -1,6 +1,11 @@
 import { QueryClient, QueryClientProvider } from '@tanstack/react-query'
 import { getDefaultConfig } from 'connectkit'
-import { useEffect, useState, type PropsWithChildren } from 'react'
+import {
+  useEffect,
+  useState,
+  type PropsWithChildren,
+  type ReactNode,
+} from 'react'
 import { createConfig, injected, WagmiProvider, type Config } from 'wagmi'
 import {
   arbitrum,
@@ -19,13 +24,16 @@ const queryClient = new QueryClient()
 const WALLETCONNECT_PROJECT_ID = '0f8a5e2cf60430a26274b421418e8a27'
 const isServer = typeof document === 'undefined'
 
-export type WalletProviderProps = PropsWithChildren
+export type WalletProviderProps = PropsWithChildren<{ fallback?: ReactNode }>
 
-export const WalletProvider = ({ children }: WalletProviderProps) => {
+export const WalletProvider = ({
+  children,
+  fallback = null,
+}: WalletProviderProps) => {
   const config = useConfig()
 
   if (config == null) {
-    return null
+    return fallback
   }
 
   return (

deployables/app/app/components/wallet/index.ts

@@ -17,3 +17,5 @@ export const WalletProvider: ComponentType<WalletProviderProps> = lazy(
     return { default: WalletProvider }
   },
 )
+
+export { ConnectWalletFallback } from './ConnectWalletFallback'

deployables/app/app/routes.ts

@@ -7,6 +7,7 @@ import {
 
 export default [
   index('routes/index.tsx'),
+  route('/connect', 'routes/connect.tsx'),
   route('/new-route', 'routes/new-route.ts'),
   route('/edit-route/:data', 'routes/edit-route.$data.tsx'),
   ...prefix('/:account/:chainId', [

deployables/app/app/routes/connect.tsx

@@ -0,0 +1,14 @@
+import { PilotType, ZodiacOsPlain } from '@zodiac/ui'
+
+const Connect = () => {
+  return (
+    <div className="flex h-full items-center justify-center">
+      <div className="flex items-center gap-4">
+        <ZodiacOsPlain className="h-6 lg:h-8" />
+        <PilotType className="h-8 lg:h-10 dark:invert" />
+      </div>
+    </div>
+  )
+}
+
+export default Connect

deployables/app/app/routes/edit-route.$data.spec.ts

@@ -27,11 +27,12 @@
   randomPrefixedAddress,
 } from '@zodiac/test-utils'
 import { formatPrefixedAddress, type ChainId } from 'ser-kit'
-import { describe, expect, it, vi } from 'vitest'
+import { beforeEach, describe, expect, it, vi } from 'vitest'
 
 const { mockGetSafesByOwner } = vi.hoisted(() => ({
-  mockGetSafesByOwner:
-    vi.fn<ReturnType<typeof initSafeApiKit>['getSafesByOwner']>(),
+  mockGetSafesByOwner: vi
+    .fn<ReturnType<typeof initSafeApiKit>['getSafesByOwner']>()
+    .mockResolvedValue({ safes: [] }),
 }))
 
 vi.mock('@zodiac/safe', async (importOriginal) => {
@@ -79,6 +80,10 @@
 const mockDryRun = vi.mocked(dryRun)
 
 describe('Edit route', () => {
+  beforeEach(() => {
+    mockFetchZodiacModules.mockResolvedValue([])
+  })
+
   describe('Label', () => {
     it('shows the name of a route', async () => {
       const route = createMockExecutionRoute({ label: 'Test route' })
@@ -333,7 +338,7 @@
       await render(`/edit-route/${btoa(JSON.stringify(route))}`)
 
       await userEvent.click(
-        screen.getByRole('combobox', { name: 'Zodiac Mod' }),
+        await screen.findByRole('combobox', { name: 'Zodiac Mod' }),
       )
       await userEvent.click(screen.getByRole('option', { name: 'Roles v2' }))
 
@@ -403,7 +408,7 @@
         await render(`/edit-route/${btoa(JSON.stringify(route))}`)
 
         await userEvent.click(
-          screen.getByRole('combobox', { name: 'Zodiac Mod' }),
+          await screen.findByRole('combobox', { name: 'Zodiac Mod' }),
         )
         await userEvent.click(screen.getByRole('option', { name: 'Roles v1' }))
 
@@ -468,7 +473,7 @@
 
           await userEvent.click(screen.getByRole('button', { name: 'Save' }))
 
           expect(mockPostMessage).toHaveBeenCalledWith(
             {
               type: CompanionAppMessageType.SAVE_ROUTE,
               data: updateRoleId(route, roleId),
@@ -520,7 +525,7 @@
         await render(`/edit-route/${btoa(JSON.stringify(route))}`)
 
         await userEvent.click(
-          screen.getByRole('combobox', { name: 'Zodiac Mod' }),
+          await screen.findByRole('combobox', { name: 'Zodiac Mod' }),
         )
         await userEvent.click(screen.getByRole('option', { name: 'Roles v2' }))
 
@@ -586,7 +591,7 @@
 
           await userEvent.click(screen.getByRole('button', { name: 'Save' }))
 
           expect(mockPostMessage).toHaveBeenCalledWith(
             {
               type: CompanionAppMessageType.SAVE_ROUTE,
               data: updateRoleId(route, encodeRoleKey('MANAGER')),

deployables/app/app/routes/edit-route.$data.tsx

@@ -2,11 +2,12 @@ import {
   AvatarInput,
   ChainSelect,
   ConnectWallet,
+  ConnectWalletFallback,
   WalletProvider,
   ZodiacMod,
 } from '@/components'
 import { dryRun, editRoute, jsonRpcProvider, parseRouteData } from '@/utils'
-import { invariant } from '@epic-web/invariant'
+import { invariant, invariantResponse } from '@epic-web/invariant'
 import { Chain, getChainId, verifyChainId, ZERO_ADDRESS } from '@zodiac/chains'
 import {
   formData,
@@ -43,12 +44,11 @@ import {
   PilotType,
   PrimaryButton,
   SecondaryButton,
-  Section,
   TextInput,
   ZodiacOsPlain,
 } from '@zodiac/ui'
-import { lazy, Suspense, useState } from 'react'
-import { Form, useNavigation, useSubmit } from 'react-router'
+import { lazy, Suspense, useEffect, useState } from 'react'
+import { Form, useLoaderData, useNavigation, useSubmit } from 'react-router'
 import type { Route } from './+types/edit-route.$data'
 import { Intent } from './intents'
 
@@ -80,6 +80,13 @@ export const action = async ({ request, params }: Route.ActionArgs) => {
   const route = parseRouteData(params.data)
   const data = await request.formData()
 
+  const intent = getString(data, 'intent')
+
+  invariantResponse(
+    intent === Intent.UpdateModule,
+    `Invalid intent "${intent}" received in server action`,
+  )
+
   const module = zodiacModuleSchema.parse(JSON.parse(getString(data, 'module')))
 
   const updatedRoute = updateRolesWaypoint(route, {
@@ -169,21 +176,16 @@ export const clientAction = async ({
 }
 
 const EditRoute = ({
-  loaderData: { chainId, label, avatar, providerType, waypoints },
+  loaderData: { chainId, label, avatar, waypoints },
   actionData,
 }: Route.ComponentProps) => {
   const submit = useSubmit()
-
-  const [optimisticConnection, setOptimisticConnection] = useState({
-    pilotAddress: getPilotAddress(waypoints),
-    chainId,
-    providerType,
-  })
+  const optimisticRoute = useOptimisticRoute()
 
   const { state } = useNavigation()
 
   return (
-    <div className="flex h-full flex-col gap-4">
+    <div className="flex h-full flex-col gap-8">
       <main className="mx-auto flex w-3/4 flex-col gap-4 md:w-1/2 2xl:w-1/4">
         <div className="flex items-center justify-between">
           <div className="flex items-center gap-4">
@@ -197,25 +199,13 @@ const EditRoute = ({
         <Form method="POST" className="flex flex-col gap-4">
           <TextInput label="Label" name="label" defaultValue={label} />
 
-          <Suspense
-            fallback={
-              <Section title="Pilot Account">
-                <PrimaryButton disabled>Connect wallet</PrimaryButton>
-              </Section>
-            }
-          >
-            <WalletProvider>
+          <Suspense fallback={<ConnectWalletFallback />}>
+            <WalletProvider fallback={<ConnectWalletFallback />}>
               <ConnectWallet
-                chainId={optimisticConnection.chainId}
-                pilotAddress={optimisticConnection.pilotAddress}
-                providerType={optimisticConnection.providerType}
+                chainId={optimisticRoute.chainId}
+                pilotAddress={optimisticRoute.pilotAddress}
+                providerType={optimisticRoute.providerType}
                 onConnect={({ account, chainId, providerType }) => {
-                  setOptimisticConnection({
-                    pilotAddress: account,
-                    chainId,
-                    providerType,
-                  })
-
                   submit(
                     formData({
                       intent: Intent.ConnectWallet,
@@ -227,12 +217,6 @@ const EditRoute = ({
                   )
                 }}
                 onDisconnect={() => {
-                  setOptimisticConnection({
-                    pilotAddress: ZERO_ADDRESS,
-                    chainId: Chain.ETH,
-                    providerType: undefined,
-                  })
-
                   submit(formData({ intent: Intent.DisconnectWallet }), {
                     method: 'POST',
                   })
@@ -270,9 +254,15 @@ const EditRoute = ({
             avatar={avatar}
             waypoints={waypoints}
             onSelect={(module) => {
-              submit(formData({ module: JSON.stringify(module) }), {
-                method: 'POST',
-              })
+              submit(
+                formData({
+                  intent: Intent.UpdateModule,
+                  module: JSON.stringify(module),
+                }),
+                {
+                  method: 'POST',
+                },
+              )
             }}
           />
 
@@ -315,6 +305,57 @@ const EditRoute = ({
 
 export default EditRoute
 
+const useOptimisticRoute = () => {
+  const { waypoints, chainId, providerType } = useLoaderData<typeof loader>()
+  const pilotAddress = getPilotAddress(waypoints)
+
+  const { formData } = useNavigation()
+
+  const [optimisticConnection, setOptimisticConnection] = useState({
+    pilotAddress,
+    chainId,
+    providerType,
+  })
+
+  useEffect(() => {
+    setOptimisticConnection({ pilotAddress, chainId, providerType })
+  }, [chainId, pilotAddress, providerType])
+
+  useEffect(() => {
+    if (formData == null) {
+      return
+    }
+
+    const intent = getOptionalString(formData, 'intent')
+
+    if (intent == null) {
+      return
+    }
+
+    switch (intent) {
+      case Intent.DisconnectWallet: {
+        setOptimisticConnection({
+          pilotAddress: ZERO_ADDRESS,
+          chainId: Chain.ETH,
+          providerType: undefined,
+        })
+
+        break
+      }
+
+      case Intent.ConnectWallet: {
+        setOptimisticConnection({
+          pilotAddress: getHexString(formData, 'account'),
+          chainId: verifyChainId(getInt(formData, 'chainId')),
+          providerType: verifyProviderType(getInt(formData, 'providerType')),
+        })
+      }
+    }
+  }, [formData])
+
+  return optimisticConnection
+}
+
 const getPilotAddress = (waypoints?: Waypoints) => {
   if (waypoints == null) {
     return null

deployables/app/app/routes/intents.ts

@@ -1,5 +1,6 @@
 export enum Intent {
   Save = 'Save',
+  UpdateModule = 'UpdateModule',
   UpdateChain = 'UpdateChain',
   UpdateAvatar = 'UpdateAvatar',
   RemoveAvatar = 'RemoveAvatar',

deployables/app/e2e/utils/index.ts

@@ -1,4 +1,2 @@
 export { expect, test } from './fixture'
 export { loadExtension } from './loadExtension'
-export { defaultMockAccount, mockWeb3 } from './mockWeb3'
-export { waitFor } from './waitFor'

deployables/app/tsconfig.cloudflare.json

@@ -3,6 +3,7 @@
   "include": [
     ".react-router/types/**/*",
     "vitest.setup.ts",
+    "e2e/**/*",
     "app/**/*",
     "app/**/.server/**/*",
     "app/**/.client/**/*",

deployables/extension/manifest-util.ts

@@ -13,10 +13,6 @@ config()
 //
 // node manifest-util.js ./public/manifest.json
 
-const getIframeUrl = () => {
-  return `${getCompanionAppUrl()}/`
-}
-
 const updateManifest = (templateFileName, outFileName, version) => {
   try {
     console.log(chalk.white.bold('Manifest template file:'))
@@ -25,7 +21,7 @@ const updateManifest = (templateFileName, outFileName, version) => {
     const data = fs
       .readFileSync(templateFileName)
       .toString()
-      .replaceAll('<COMPANION_APP_URL>', getIframeUrl())
+      .replaceAll('<COMPANION_APP_URL>', getCompanionAppUrl())
 
     const manifest = JSON.parse(data)
     manifest['version'] = version.replace('v', '')

deployables/extension/src/connect/contentScripts/dApp.ts

@@ -5,16 +5,16 @@ import {
   type ConnectedWalletMessage,
 } from '@zodiac/messages'
 
-const companionAppUrl = getCompanionAppUrl()
+const connectUrl = `${getCompanionAppUrl()}/connect`
 
 const ensureIframe = () => {
   let node: HTMLIFrameElement | null = document.querySelector(
-    `iframe[src="${companionAppUrl}"]`,
+    `iframe[src="${connectUrl}"]`,
   )
 
   if (!node) {
     node = document.createElement('iframe')
-    node.src = companionAppUrl
+    node.src = connectUrl
     node.style.display = 'none'
 
     const parent = document.body || document.documentElement
@@ -40,7 +40,7 @@ chrome.runtime.onConnect.addListener((port) => {
       'cannot access connect iframe window',
     )
 
-    iframe.contentWindow.postMessage(message, companionAppUrl)
+    iframe.contentWindow.postMessage(message, connectUrl)
 
     // wait for response
     const handleResponse = (event: MessageEvent<ConnectedWalletMessage>) => {