refactor(ui): extract useMutation (#1010)

ee/tabby-ui/app/(dashboard)/cluster/components/cluster.tsx

@@ -4,7 +4,7 @@ import { graphql } from '@/lib/gql/generates'
 import { WorkerKind } from '@/lib/gql/generates/graphql'
 import { useHealth } from '@/lib/hooks/use-health'
 import { useWorkers } from '@/lib/hooks/use-workers'
-import { useAuthenticatedGraphQLQuery, useGraphQLForm } from '@/lib/tabby/gql'
+import { useAuthenticatedGraphQLQuery, useMutation } from '@/lib/tabby/gql'
 import { Button } from '@/components/ui/button'
 import { IconRotate } from '@/components/ui/icons'
 import { Input } from '@/components/ui/input'
@@ -36,12 +36,11 @@ export default function Workers() {
     getRegistrationTokenDocument
   )
 
-  const { onSubmit: resetRegistrationToken } = useGraphQLForm(
-    resetRegistrationTokenDocument,
-    {
-      onSuccess: () => mutate()
+  const resetRegistrationToken = useMutation(resetRegistrationTokenDocument, {
+    onCompleted() {
+      mutate()
     }
-  )
+  })
 
   if (!healthInfo) return
 

ee/tabby-ui/app/(dashboard)/page.tsx

@@ -4,7 +4,7 @@ import { useEffect, useState } from 'react'
 
 import { graphql } from '@/lib/gql/generates'
 import { useHealth } from '@/lib/hooks/use-health'
-import { useAuthenticatedGraphQLQuery, useGraphQLForm } from '@/lib/tabby/gql'
+import { useAuthenticatedGraphQLQuery, useMutation } from '@/lib/tabby/gql'
 import { Button } from '@/components/ui/button'
 import {
   CardContent,
@@ -49,12 +49,9 @@ function MainPanel() {
     setOrigin(new URL(window.location.href).origin)
   }, [])
 
-  const { onSubmit: resetUserAuthToken } = useGraphQLForm(
-    resetUserAuthTokenDocument,
-    {
-      onSuccess: () => mutate()
-    }
-  )
+  const resetUserAuthToken = useMutation(resetUserAuthTokenDocument, {
+    onCompleted: () => mutate()
+  })
 
   if (!healthInfo || !data) return
 

ee/tabby-ui/app/(dashboard)/team/components/create-invitation-form.tsx

@@ -6,7 +6,7 @@ import { useForm } from 'react-hook-form'
 import * as z from 'zod'
 
 import { graphql } from '@/lib/gql/generates'
-import { useGraphQLForm } from '@/lib/tabby/gql'
+import { useMutation } from '@/lib/tabby/gql'
 import { Button } from '@/components/ui/button'
 import {
   Form,
@@ -17,7 +17,7 @@ import {
 } from '@/components/ui/form'
 import { Input } from '@/components/ui/input'
 
-const createInvitation = graphql(/* GraphQL */ `
+const createInvitationMutation = graphql(/* GraphQL */ `
   mutation CreateInvitation($email: String!) {
     createInvitation(email: $email)
   }
@@ -37,12 +37,12 @@ export default function CreateInvitationForm({
   })
 
   const { isSubmitting } = form.formState
-  const { onSubmit } = useGraphQLForm(createInvitation, {
-    onSuccess: () => {
+  const createInvitation = useMutation(createInvitationMutation, {
+    onCompleted() {
       form.reset({ email: '' })
       onCreated()
     },
-    onError: (path, message) => form.setError(path as any, { message })
+    form
   })
 
   return (
@@ -50,7 +50,7 @@ export default function CreateInvitationForm({
       <div className="flex flex-col items-start gap-2">
         <form
           className="flex w-full items-center gap-2"
-          onSubmit={form.handleSubmit(onSubmit)}
+          onSubmit={form.handleSubmit(createInvitation)}
         >
           <FormField
             control={form.control}

ee/tabby-ui/app/(dashboard)/team/components/invitation-table.tsx

@@ -4,7 +4,7 @@ import React, { useEffect, useState } from 'react'
 import moment from 'moment'
 
 import { graphql } from '@/lib/gql/generates'
-import { useAuthenticatedGraphQLQuery, useGraphQLForm } from '@/lib/tabby/gql'
+import { useAuthenticatedGraphQLQuery, useMutation } from '@/lib/tabby/gql'
 import { Button } from '@/components/ui/button'
 import { IconTrash } from '@/components/ui/icons'
 import {
@@ -44,12 +44,11 @@ export default function InvitationTable() {
     setOrigin(new URL(window.location.href).origin)
   }, [])
 
-  const { onSubmit: deleteInvitation } = useGraphQLForm(
-    deleteInvitationMutation,
-    {
-      onSuccess: () => mutate()
+  const deleteInvitation = useMutation(deleteInvitationMutation, {
+    onCompleted() {
+      mutate()
     }
-  )
+  })
 
   return (
     invitations && (
@@ -75,7 +74,7 @@ export default function InvitationTable() {
                   <Button
                     size="icon"
                     variant="hover-destructive"
-                    onClick={() => deleteInvitation({ id: x.id })}
+                    onClick={() => deleteInvitation(x)}
                   >
                     <IconTrash />
                   </Button>

ee/tabby-ui/app/auth/signin/components/user-signin-form.tsx

@@ -8,7 +8,7 @@ import * as z from 'zod'
 
 import { graphql } from '@/lib/gql/generates'
 import { useSignIn } from '@/lib/tabby/auth'
-import { useGraphQLForm } from '@/lib/tabby/gql'
+import { useMutation } from '@/lib/tabby/gql'
 import { cn } from '@/lib/utils'
 import { Button } from '@/components/ui/button'
 import {
@@ -52,13 +52,13 @@ export default function UserSignInForm({
   const router = useRouter()
   const signIn = useSignIn()
   const { isSubmitting } = form.formState
-  const { onSubmit } = useGraphQLForm(tokenAuth, {
-    onSuccess: async values => {
+  const onSubmit = useMutation(tokenAuth, {
+    async onCompleted(values) {
       if (await signIn(values.tokenAuth)) {
         router.replace('/')
       }
     },
-    onError: (path, message) => form.setError(path as any, { message })
+    form
   })
 
   return (

ee/tabby-ui/app/auth/signup/components/user-register-form.tsx

@@ -8,7 +8,7 @@ import * as z from 'zod'
 
 import { graphql } from '@/lib/gql/generates'
 import { useSignIn } from '@/lib/tabby/auth'
-import { useGraphQLForm } from '@/lib/tabby/gql'
+import { useMutation } from '@/lib/tabby/gql'
 import { cn } from '@/lib/utils'
 import { Button } from '@/components/ui/button'
 import {
@@ -67,13 +67,13 @@ export function UserAuthForm({
   const router = useRouter()
   const signIn = useSignIn()
   const { isSubmitting } = form.formState
-  const { onSubmit } = useGraphQLForm(registerUser, {
-    onSuccess: async values => {
+  const onSubmit = useMutation(registerUser, {
+    async onCompleted(values) {
       if (await signIn(values.register)) {
         router.replace('/')
       }
     },
-    onError: (path, message) => form.setError(path as any, { message })
+    form
   })
 
   return (

ee/tabby-ui/lib/tabby/auth.tsx

@@ -4,7 +4,7 @@ import { jwtDecode, JwtPayload } from 'jwt-decode'
 
 import { graphql } from '@/lib/gql/generates'
 import useInterval from '@/lib/hooks/use-interval'
-import { gqlClient, useGraphQLQuery } from '@/lib/tabby/gql'
+import { useGraphQLQuery, useMutation } from '@/lib/tabby/gql'
 
 interface AuthData {
   accessToken: string
@@ -115,35 +115,36 @@ const refreshTokenMutation = graphql(/* GraphQL */ `
   }
 `)
 
-async function doRefresh(token: string, dispatch: React.Dispatch<AuthActions>) {
-  let action: AuthActions
-  try {
-    action = {
-      type: AuthActionType.Refresh,
-      data: (
-        await gqlClient.request(refreshTokenMutation, { refreshToken: token })
-      ).refreshToken
-    }
-  } catch (err) {
-    console.error('Failed to refresh token', err)
-    action = {
-      type: AuthActionType.SignOut
-    }
-  }
-
-  dispatch(action)
-}
-
 const AuthProvider: React.FunctionComponent<AuthProviderProps> = ({
   children
 }) => {
-  const storage = new TokenStorage()
-
   const [authState, dispatch] = React.useReducer(authReducer, {
     status: 'loading',
     data: null
   })
 
+  return (
+    <AuthContext.Provider value={{ authState, dispatch }}>
+      <RefreshAuth />
+      {children}
+    </AuthContext.Provider>
+  )
+}
+
+function RefreshAuth() {
+  const { authState, dispatch } = useAuthStore()
+  const storage = new TokenStorage()
+  const refreshToken = useMutation(refreshTokenMutation, {
+    onCompleted({ refreshToken: data }) {
+      dispatch({ type: AuthActionType.Refresh, data })
+    },
+    onError() {
+      dispatch({
+        type: AuthActionType.SignOut
+      })
+    }
+  })
+
   const initialized = React.useRef(false)
   React.useEffect(() => {
     if (initialized.current) return
@@ -151,38 +152,38 @@ const AuthProvider: React.FunctionComponent<AuthProviderProps> = ({
     initialized.current = true
     const data = storage.initialState()
     if (data?.refreshToken) {
-      doRefresh(data.refreshToken, dispatch)
+      refreshToken(data)
     } else {
       dispatch({ type: AuthActionType.Init, data: null })
     }
   }, [])
 
   React.useEffect(() => {
-    authState.data && storage.persist(authState.data)
+    authState?.data && storage.persist(authState.data)
   }, [authState])
 
   useInterval(async () => {
-    if (authState.status !== 'authenticated') {
+    if (authState?.status !== 'authenticated') {
       return
     }
 
-    await doRefresh(authState.data.refreshToken, dispatch)
+    await refreshToken(authState.data)
   }, 5)
 
-  return (
-    <AuthContext.Provider value={{ authState, dispatch }}>
-      {children}
-    </AuthContext.Provider>
-  )
+  return <></>
+}
+
+class AuthProviderIsMissing extends Error {
+  constructor() {
+    super('AuthProvider is missing. Please add the AuthProvider at root level')
+  }
 }
 
 function useAuthStore(): AuthStore {
   const context = React.useContext(AuthContext)
 
   if (!context) {
-    throw new Error(
-      'AuthProvider is missing. Please add the AuthProvider at root level'
-    )
+    throw new AuthProviderIsMissing()
   }
 
   return context

ee/tabby-ui/lib/tabby/gql.ts

@@ -1,11 +1,12 @@
 import { TypedDocumentNode } from '@graphql-typed-document-node/core'
 import { GraphQLClient, Variables } from 'graphql-request'
 import { GraphQLResponse } from 'graphql-request/build/esm/types'
+import { FieldValues, UseFormReturn } from 'react-hook-form'
 import useSWR, { SWRConfiguration, SWRResponse } from 'swr'
 
 import { useSession } from './auth'
 
-export const gqlClient = new GraphQLClient(
+const gqlClient = new GraphQLClient(
   `${process.env.NEXT_PUBLIC_TABBY_SERVER_URL ?? ''}/graphql`
 )
 
@@ -18,52 +19,59 @@ export interface ValidationErrors {
   errors: Array<ValidationError>
 }
 
-export function useGraphQLForm<
-  TResult,
-  TVariables extends Variables | undefined
->(
+export function useMutation<TResult, TVariables extends Variables | undefined>(
   document: TypedDocumentNode<TResult, TVariables>,
   options?: {
-    onSuccess?: (values: TResult) => void
-    onError?: (path: string, message: string) => void
+    onCompleted?: (data: TResult) => void
+    onError?: (err: any) => any
+    form?: any
   }
 ) {
-  const { data } = useSession()
-  const accessToken = data?.accessToken
-  const onSubmit = async (variables?: TVariables) => {
-    let res
+  const { data: session } = useSession()
+  const onFormError = options?.form
+    ? makeFormErrorHandler(options.form)
+    : undefined
+
+  const fn = async (variables?: TVariables) => {
+    let res: TResult | undefined
     try {
       res = await gqlClient.request({
         document,
-        variables,
-        requestHeaders: accessToken
+        variables: variables,
+        requestHeaders: session
           ? {
-              authorization: `Bearer ${accessToken}`
+              authorization: `Bearer ${session.accessToken}`
             }
           : undefined
       })
     } catch (err) {
-      console.error('err', err)
-      const { errors = [] } = (err as any).response as GraphQLResponse
-      for (const error of errors) {
-        if (error.extensions && error.extensions['validation-errors']) {
-          const validationErrors = error.extensions[
-            'validation-errors'
-          ] as ValidationErrors
-          for (const error of validationErrors.errors) {
-            options?.onError && options?.onError(error.path, error.message)
-          }
-        } else {
-          options?.onError && options?.onError('root', error.message)
-        }
-      }
-
-      return res
+      onFormError && onFormError(err)
+      options?.onError && options.onError(err)
+      return
     }
 
-    options?.onSuccess && options.onSuccess(res)
+    options?.onCompleted && options.onCompleted(res)
+  }
+
+  return fn
+}
+
+function makeFormErrorHandler<T extends FieldValues>(form: UseFormReturn<T>) {
+  return (err: any) => {
+    const { errors = [] } = err.response as GraphQLResponse
+    for (const error of errors) {
+      if (error.extensions && error.extensions['validation-errors']) {
+        const validationErrors = error.extensions[
+          'validation-errors'
+        ] as ValidationErrors
+        for (const error of validationErrors.errors) {
+          form.setError(error.path as any, error)
+        }
+      } else {
+        form.setError('root', error)
+      }
+    }
   }
-  return { onSubmit }
 }
 
 export function useGraphQLQuery<