From 779eef0c1f3b112b6af0ceef596e9102428d2efd Mon Sep 17 00:00:00 2001
From: Meng Zhang <meng@tabbyml.com>
Date: Thu, 7 Dec 2023 17:56:08 +0800
Subject: [PATCH] feat: add register / signin / signup / logout workflow (#976)

* feat: extract signup page

* feat: add signin page

* feat: implement auto redirecting

* fix: show slack dialog after signin

* feat: tweak signin page

* feat: add basic user panel + signout

* [autofix.ci] apply automated fixes

* fix: correct page title

* fix lint

---------

Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
---
 .../app/(dashboard)/components/sidebar.tsx    |   7 +-
 ee/tabby-ui/app/(dashboard)/page.tsx          |   6 +-
 ee/tabby-ui/app/auth/components/signup.tsx    |  34 ------
 ee/tabby-ui/app/auth/layout.tsx               |  11 ++
 ee/tabby-ui/app/auth/page.tsx                 |  16 ---
 .../app/auth/signin/components/signin.tsx     |  17 +++
 .../signin/components/user-signin-form.tsx    | 113 ++++++++++++++++++
 ee/tabby-ui/app/auth/signin/page.tsx          |  10 ++
 .../app/auth/signup/components/signup.tsx     |  50 ++++++++
 .../components/user-register-form.tsx}        |  16 ++-
 ee/tabby-ui/app/auth/signup/page.tsx          |  10 ++
 .../playground/components/chat-sessions.tsx   |   1 +
 ee/tabby-ui/components/header.tsx             |  33 +++++
 ee/tabby-ui/components/ui/icons.tsx           |  25 +++-
 ee/tabby-ui/components/user-panel.tsx         |  20 ++++
 ee/tabby-ui/lib/gql/generates/gql.ts          |  16 ++-
 ee/tabby-ui/lib/gql/generates/graphql.ts      | 108 +++++++++++++++--
 ee/tabby-ui/lib/tabby/auth.tsx                |  13 +-
 ee/tabby-ui/lib/tabby/gql.ts                  |   8 +-
 19 files changed, 436 insertions(+), 78 deletions(-)
 delete mode 100644 ee/tabby-ui/app/auth/components/signup.tsx
 create mode 100644 ee/tabby-ui/app/auth/layout.tsx
 delete mode 100644 ee/tabby-ui/app/auth/page.tsx
 create mode 100644 ee/tabby-ui/app/auth/signin/components/signin.tsx
 create mode 100644 ee/tabby-ui/app/auth/signin/components/user-signin-form.tsx
 create mode 100644 ee/tabby-ui/app/auth/signin/page.tsx
 create mode 100644 ee/tabby-ui/app/auth/signup/components/signup.tsx
 rename ee/tabby-ui/app/auth/{components/user-auth-form.tsx => signup/components/user-register-form.tsx} (90%)
 create mode 100644 ee/tabby-ui/app/auth/signup/page.tsx
 create mode 100644 ee/tabby-ui/components/user-panel.tsx

diff --git a/ee/tabby-ui/app/(dashboard)/components/sidebar.tsx b/ee/tabby-ui/app/(dashboard)/components/sidebar.tsx
index 310360f..ed273ae 100644
--- a/ee/tabby-ui/app/(dashboard)/components/sidebar.tsx
+++ b/ee/tabby-ui/app/(dashboard)/components/sidebar.tsx
@@ -1,5 +1,6 @@
 'use client'
 
+import UserPanel from '@/components/user-panel'
 import { cn } from '@/lib/utils'
 import { cva } from 'class-variance-authority'
 import Link from 'next/link'
@@ -16,7 +17,7 @@ export default function Sidebar({ children, className }: SidebarProps) {
       className={cn('grid overflow-hidden lg:grid-cols-[280px_1fr]', className)}
     >
       <div className="hidden border-r bg-zinc-100/40 dark:bg-zinc-800/40 lg:block">
-        <div className="flex flex-col gap-2">
+        <div className="flex flex-col gap-2 h-full">
           <div className="h-[12px]"></div>
           <div className="flex-1">
             <nav className="grid items-start gap-4 px-4 text-sm font-medium">
@@ -62,6 +63,10 @@ export default function Sidebar({ children, className }: SidebarProps) {
               </SidebarButton>
             </nav>
           </div>
+
+          <div className="mt-auto">
+            <UserPanel />
+          </div>
         </div>
       </div>
       <div className="flex flex-1 flex-col overflow-auto">{children}</div>
diff --git a/ee/tabby-ui/app/(dashboard)/page.tsx b/ee/tabby-ui/app/(dashboard)/page.tsx
index 00c7ff3..fdd777a 100644
--- a/ee/tabby-ui/app/(dashboard)/page.tsx
+++ b/ee/tabby-ui/app/(dashboard)/page.tsx
@@ -19,17 +19,21 @@ import { WorkerKind } from '@/lib/gql/generates/graphql'
 import { CopyButton } from '@/components/copy-button'
 import { graphql } from '@/lib/gql/generates'
 import { useGraphQLQuery } from '@/lib/tabby/gql'
+import { useSession } from '@/lib/tabby/auth'
 
 const COMMUNITY_DIALOG_SHOWN_KEY = 'community-dialog-shown'
 
 export default function Home() {
+  const { status } = useSession()
   const [open, setOpen] = useState(false)
   useEffect(() => {
+    if (status !== 'authenticated') return
+
     if (!localStorage.getItem(COMMUNITY_DIALOG_SHOWN_KEY)) {
       setOpen(true)
       localStorage.setItem(COMMUNITY_DIALOG_SHOWN_KEY, 'true')
     }
-  }, [])
+  }, [status])
 
   return (
     <div className="p-4 lg:p-16">
diff --git a/ee/tabby-ui/app/auth/components/signup.tsx b/ee/tabby-ui/app/auth/components/signup.tsx
deleted file mode 100644
index d8ab4e6..0000000
--- a/ee/tabby-ui/app/auth/components/signup.tsx
+++ /dev/null
@@ -1,34 +0,0 @@
-'use client'
-
-import { graphql } from '@/lib/gql/generates'
-import { UserAuthForm } from './user-auth-form'
-import { useSearchParams } from 'next/navigation'
-import { useGraphQLQuery } from '@/lib/tabby/gql'
-
-export const getIsAdminInitialized = graphql(/* GraphQL */ `
-  query GetIsAdminInitialized {
-    isAdminInitialized
-  }
-`)
-
-export default function Signup() {
-  const { data } = useGraphQLQuery(getIsAdminInitialized)
-  const title = data?.isAdminInitialized
-    ? 'Create an account'
-    : 'Create an admin account'
-
-  const searchParams = useSearchParams()
-  const invitationCode = searchParams.get('invitationCode') || undefined
-
-  return (
-    <div className="space-y-6 w-[350px]">
-      <div className="flex flex-col space-y-2 text-center">
-        <h1 className="text-2xl font-semibold tracking-tight">{title}</h1>
-        <p className="text-sm text-muted-foreground">
-          Fill form below to create your account
-        </p>
-      </div>
-      <UserAuthForm invitationCode={invitationCode} />
-    </div>
-  )
-}
diff --git a/ee/tabby-ui/app/auth/layout.tsx b/ee/tabby-ui/app/auth/layout.tsx
new file mode 100644
index 0000000..e8cae81
--- /dev/null
+++ b/ee/tabby-ui/app/auth/layout.tsx
@@ -0,0 +1,11 @@
+export default function RootLayout({
+  children
+}: {
+  children: React.ReactNode
+}) {
+  return (
+    <div className="flex flex-col items-center justify-center flex-1">
+      {children}
+    </div>
+  )
+}
diff --git a/ee/tabby-ui/app/auth/page.tsx b/ee/tabby-ui/app/auth/page.tsx
deleted file mode 100644
index f6a9db4..0000000
--- a/ee/tabby-ui/app/auth/page.tsx
+++ /dev/null
@@ -1,16 +0,0 @@
-import { Metadata } from 'next'
-
-import Signup from './components/signup'
-
-export const metadata: Metadata = {
-  title: 'Authentication',
-  description: 'Authentication forms built using the components.'
-}
-
-export default function AuthenticationPage() {
-  return (
-    <div className="flex flex-col items-center justify-center flex-1">
-      <Signup />
-    </div>
-  )
-}
diff --git a/ee/tabby-ui/app/auth/signin/components/signin.tsx b/ee/tabby-ui/app/auth/signin/components/signin.tsx
new file mode 100644
index 0000000..1fc0929
--- /dev/null
+++ b/ee/tabby-ui/app/auth/signin/components/signin.tsx
@@ -0,0 +1,17 @@
+'use client'
+
+import UserSignInForm from './user-signin-form'
+
+export default function Signin() {
+  return (
+    <div className="space-y-6 w-[350px]">
+      <div className="flex flex-col space-y-2 text-center">
+        <h1 className="text-2xl font-semibold tracking-tight">Sign In</h1>
+        <p className="text-sm text-muted-foreground">
+          Enter credentials to login to your account
+        </p>
+      </div>
+      <UserSignInForm />
+    </div>
+  )
+}
diff --git a/ee/tabby-ui/app/auth/signin/components/user-signin-form.tsx b/ee/tabby-ui/app/auth/signin/components/user-signin-form.tsx
new file mode 100644
index 0000000..fb85323
--- /dev/null
+++ b/ee/tabby-ui/app/auth/signin/components/user-signin-form.tsx
@@ -0,0 +1,113 @@
+'use client'
+
+import * as React from 'react'
+
+import { zodResolver } from '@hookform/resolvers/zod'
+import { useForm } from 'react-hook-form'
+import * as z from 'zod'
+
+import { cn } from '@/lib/utils'
+import { IconSpinner } from '@/components/ui/icons'
+import { Button } from '@/components/ui/button'
+import { Input } from '@/components/ui/input'
+import {
+  Form,
+  FormControl,
+  FormField,
+  FormItem,
+  FormLabel,
+  FormMessage
+} from '@/components/ui/form'
+import { graphql } from '@/lib/gql/generates'
+import { useGraphQLForm } from '@/lib/tabby/gql'
+import { useSignIn } from '@/lib/tabby/auth'
+import { useRouter } from 'next/navigation'
+
+export const tokenAuth = graphql(/* GraphQL */ `
+  mutation tokenAuth($email: String!, $password: String!) {
+    tokenAuth(email: $email, password: $password) {
+      accessToken
+      refreshToken
+    }
+  }
+`)
+
+const formSchema = z.object({
+  email: z.string().email('Invalid email address'),
+  password: z.string()
+})
+
+interface UserAuthFormProps extends React.HTMLAttributes<HTMLDivElement> {
+  invitationCode?: string
+}
+
+export default function UserSignInForm({
+  className,
+  invitationCode,
+  ...props
+}: UserAuthFormProps) {
+  const form = useForm<z.infer<typeof formSchema>>({
+    resolver: zodResolver(formSchema)
+  })
+
+  const router = useRouter()
+  const signIn = useSignIn()
+  const { isSubmitting } = form.formState
+  const { onSubmit } = useGraphQLForm(tokenAuth, {
+    onSuccess: async values => {
+      if (await signIn(values.tokenAuth)) {
+        router.replace('/')
+      }
+    },
+    onError: (path, message) => form.setError(path as any, { message })
+  })
+
+  return (
+    <div className={cn('grid gap-6', className)} {...props}>
+      <Form {...form}>
+        <form className="grid gap-2" onSubmit={form.handleSubmit(onSubmit)}>
+          <FormField
+            control={form.control}
+            name="email"
+            render={({ field }) => (
+              <FormItem>
+                <FormLabel>Email</FormLabel>
+                <FormControl>
+                  <Input
+                    placeholder=""
+                    type="email"
+                    autoCapitalize="none"
+                    autoComplete="email"
+                    autoCorrect="off"
+                    {...field}
+                  />
+                </FormControl>
+                <FormMessage />
+              </FormItem>
+            )}
+          />
+          <FormField
+            control={form.control}
+            name="password"
+            render={({ field }) => (
+              <FormItem>
+                <FormLabel>Password</FormLabel>
+                <FormControl>
+                  <Input type="password" {...field} />
+                </FormControl>
+                <FormMessage />
+              </FormItem>
+            )}
+          />
+          <Button type="submit" className="mt-1" disabled={isSubmitting}>
+            {isSubmitting && (
+              <IconSpinner className="mr-2 h-4 w-4 animate-spin" />
+            )}
+            Login
+          </Button>
+        </form>
+        <FormMessage className="text-center" />
+      </Form>
+    </div>
+  )
+}
diff --git a/ee/tabby-ui/app/auth/signin/page.tsx b/ee/tabby-ui/app/auth/signin/page.tsx
new file mode 100644
index 0000000..07806cf
--- /dev/null
+++ b/ee/tabby-ui/app/auth/signin/page.tsx
@@ -0,0 +1,10 @@
+import { Metadata } from 'next'
+import Signin from './components/signin'
+
+export const metadata: Metadata = {
+  title: 'Sign In'
+}
+
+export default function Page() {
+  return <Signin />
+}
diff --git a/ee/tabby-ui/app/auth/signup/components/signup.tsx b/ee/tabby-ui/app/auth/signup/components/signup.tsx
new file mode 100644
index 0000000..4b53c70
--- /dev/null
+++ b/ee/tabby-ui/app/auth/signup/components/signup.tsx
@@ -0,0 +1,50 @@
+'use client'
+
+import { UserAuthForm } from './user-register-form'
+import { useSearchParams } from 'next/navigation'
+
+export default function Signup() {
+  const searchParams = useSearchParams()
+  const invitationCode = searchParams.get('invitationCode') || undefined
+  const isAdmin = searchParams.get('isAdmin') || false
+
+  const title = isAdmin ? 'Create an admin account' : 'Create an account'
+
+  const description = isAdmin
+    ? 'The admin account has access to invite collaborators and manage Tabby configuration'
+    : 'Fill form below to create your account'
+
+  if (isAdmin || invitationCode) {
+    return <Content title={title} description={description} show />
+  } else {
+    return (
+      <Content
+        title="No invitation code"
+        description="Please contact your Tabby admin for an invitation code to register"
+      />
+    )
+  }
+}
+
+function Content({
+  title,
+  description,
+  show
+}: {
+  title: string
+  description: string
+  show?: boolean
+}) {
+  const searchParams = useSearchParams()
+  const invitationCode = searchParams.get('invitationCode') || undefined
+
+  return (
+    <div className="space-y-6 w-[350px]">
+      <div className="flex flex-col space-y-2 text-center">
+        <h1 className="text-2xl font-semibold tracking-tight">{title}</h1>
+        <p className="text-sm text-muted-foreground">{description}</p>
+      </div>
+      {show && <UserAuthForm invitationCode={invitationCode} />}
+    </div>
+  )
+}
diff --git a/ee/tabby-ui/app/auth/components/user-auth-form.tsx b/ee/tabby-ui/app/auth/signup/components/user-register-form.tsx
similarity index 90%
rename from ee/tabby-ui/app/auth/components/user-auth-form.tsx
rename to ee/tabby-ui/app/auth/signup/components/user-register-form.tsx
index 378add8..1eba9c5 100644
--- a/ee/tabby-ui/app/auth/components/user-auth-form.tsx
+++ b/ee/tabby-ui/app/auth/signup/components/user-register-form.tsx
@@ -20,6 +20,8 @@ import {
 } from '@/components/ui/form'
 import { graphql } from '@/lib/gql/generates'
 import { useGraphQLForm } from '@/lib/tabby/gql'
+import { useSignIn } from '@/lib/tabby/auth'
+import { useRouter } from 'next/navigation'
 
 export const registerUser = graphql(/* GraphQL */ `
   mutation register(
@@ -59,15 +61,19 @@ export function UserAuthForm({
   const form = useForm<z.infer<typeof formSchema>>({
     resolver: zodResolver(formSchema),
     defaultValues: {
-      email: '',
-      password1: '',
-      password2: '',
       invitationCode
     }
   })
 
+  const router = useRouter()
+  const signIn = useSignIn()
   const { isSubmitting } = form.formState
   const { onSubmit } = useGraphQLForm(registerUser, {
+    onSuccess: async values => {
+      if (await signIn(values.register)) {
+        router.replace('/')
+      }
+    },
     onError: (path, message) => form.setError(path as any, { message })
   })
 
@@ -125,14 +131,14 @@ export function UserAuthForm({
             control={form.control}
             name="invitationCode"
             render={({ field }) => (
-              <FormItem>
+              <FormItem className="hidden">
                 <FormControl>
                   <Input type="hidden" {...field} />
                 </FormControl>
               </FormItem>
             )}
           />
-          <Button type="submit" disabled={isSubmitting}>
+          <Button type="submit" className="mt-1" disabled={isSubmitting}>
             {isSubmitting && (
               <IconSpinner className="mr-2 h-4 w-4 animate-spin" />
             )}
diff --git a/ee/tabby-ui/app/auth/signup/page.tsx b/ee/tabby-ui/app/auth/signup/page.tsx
new file mode 100644
index 0000000..ccebb04
--- /dev/null
+++ b/ee/tabby-ui/app/auth/signup/page.tsx
@@ -0,0 +1,10 @@
+import { Metadata } from 'next'
+import Signup from './components/signup'
+
+export const metadata: Metadata = {
+  title: 'Sign Up'
+}
+
+export default function Page() {
+  return <Signup />
+}
diff --git a/ee/tabby-ui/app/playground/components/chat-sessions.tsx b/ee/tabby-ui/app/playground/components/chat-sessions.tsx
index 970788c..64a5e9c 100644
--- a/ee/tabby-ui/app/playground/components/chat-sessions.tsx
+++ b/ee/tabby-ui/app/playground/components/chat-sessions.tsx
@@ -20,6 +20,7 @@ import { Button } from '@/components/ui/button'
 import { ListSkeleton } from '@/components/skeleton'
 import { Separator } from '@/components/ui/separator'
 import { ClearChatsButton } from './clear-chats-button'
+import UserPanel from '@/components/user-panel'
 
 interface ChatSessionsProps {
   className?: string
diff --git a/ee/tabby-ui/components/header.tsx b/ee/tabby-ui/components/header.tsx
index 28f59f7..e0a682a 100644
--- a/ee/tabby-ui/components/header.tsx
+++ b/ee/tabby-ui/components/header.tsx
@@ -12,8 +12,14 @@ import { useWorkers } from '@/lib/hooks/use-workers'
 import { WorkerKind } from '@/lib/gql/generates/graphql'
 import { has } from 'lodash-es'
 import { ThemeToggle } from './theme-toggle'
+import { useSession } from '@/lib/tabby/auth'
+import { useRouter } from 'next/navigation'
+import { graphql } from '@/lib/gql/generates'
+import { useGraphQLQuery } from '@/lib/tabby/gql'
 
 export function Header() {
+  useRequireAuth()
+
   const { data } = useHealth()
   const workers = useWorkers(data)
   const isChatEnabled = has(workers, WorkerKind.Chat)
@@ -74,3 +80,30 @@ function isNewVersionAvailable(version?: string, latestRelease?: ReleaseInfo) {
     return true
   }
 }
+
+export const getIsAdminInitialized = graphql(/* GraphQL */ `
+  query GetIsAdminInitialized {
+    isAdminInitialized
+  }
+`)
+
+function useRequireAuth() {
+  const { data, isLoading } = useGraphQLQuery(
+    getIsAdminInitialized,
+    undefined,
+    {}
+  )
+  const router = useRouter()
+  const { status } = useSession()
+
+  React.useEffect(() => {
+    if (isLoading) return
+    if (status !== 'unauthenticated') return
+
+    if (data!.isAdminInitialized) {
+      router.replace('/auth/signin')
+    } else {
+      router.replace('/auth/signup?isAdmin=true')
+    }
+  }, [data, isLoading, status])
+}
diff --git a/ee/tabby-ui/components/ui/icons.tsx b/ee/tabby-ui/components/ui/icons.tsx
index eabc2bf..dd4de39 100644
--- a/ee/tabby-ui/components/ui/icons.tsx
+++ b/ee/tabby-ui/components/ui/icons.tsx
@@ -572,6 +572,28 @@ function IconSymbolFunction({
   )
 }
 
+function IconLogout({ className, ...props }: React.ComponentProps<'svg'>) {
+  return (
+    <svg
+      xmlns="http://www.w3.org/2000/svg"
+      width="24"
+      height="24"
+      viewBox="0 0 24 24"
+      fill="none"
+      stroke="currentColor"
+      strokeWidth="2"
+      strokeLinecap="round"
+      strokeLinejoin="round"
+      className={cn('h-4 w-4', className)}
+      {...props}
+    >
+      <path d="M9 21H5a2 2 0 0 1-2-2V5a2 2 0 0 1 2-2h4" />
+      <polyline points="16 17 21 12 16 7" />
+      <line x1="21" x2="9" y1="12" y2="12" />
+    </svg>
+  )
+}
+
 export {
   IconEdit,
   IconNextChat,
@@ -602,5 +624,6 @@ export {
   IconChevronUpDown,
   IconSlack,
   IconNotice,
-  IconSymbolFunction
+  IconSymbolFunction,
+  IconLogout
 }
diff --git a/ee/tabby-ui/components/user-panel.tsx b/ee/tabby-ui/components/user-panel.tsx
new file mode 100644
index 0000000..a5953bb
--- /dev/null
+++ b/ee/tabby-ui/components/user-panel.tsx
@@ -0,0 +1,20 @@
+import { useSession, useSignOut } from '@/lib/tabby/auth'
+import { cn } from '@/lib/utils'
+import { IconLogout } from './ui/icons'
+
+export default function UserPanel() {
+  const { data: session, status } = useSession()
+  const signOut = useSignOut()
+
+  if (status !== 'authenticated') return
+  return (
+    <div className="py-4 flex justify-center text-sm font-medium">
+      <span className={cn('flex items-center gap-2')}>
+        <span title="Sign out">
+          <IconLogout className="cursor-pointer" onClick={signOut} />
+        </span>
+        {session.email}
+      </span>
+    </div>
+  )
+}
diff --git a/ee/tabby-ui/lib/gql/generates/gql.ts b/ee/tabby-ui/lib/gql/generates/gql.ts
index e0d873e..8b0d1b5 100644
--- a/ee/tabby-ui/lib/gql/generates/gql.ts
+++ b/ee/tabby-ui/lib/gql/generates/gql.ts
@@ -15,10 +15,12 @@ import { TypedDocumentNode as DocumentNode } from '@graphql-typed-document-node/
 const documents = {
   '\n  query GetRegistrationToken {\n    registrationToken\n  }\n':
     types.GetRegistrationTokenDocument,
-  '\n  query GetIsAdminInitialized {\n    isAdminInitialized\n  }\n':
-    types.GetIsAdminInitializedDocument,
+  '\n  mutation tokenAuth($email: String!, $password: String!) {\n    tokenAuth(email: $email, password: $password) {\n      accessToken\n      refreshToken\n    }\n  }\n':
+    types.TokenAuthDocument,
   '\n  mutation register(\n    $email: String!\n    $password1: String!\n    $password2: String!\n    $invitationCode: String\n  ) {\n    register(\n      email: $email\n      password1: $password1\n      password2: $password2\n      invitationCode: $invitationCode\n    ) {\n      accessToken\n      refreshToken\n    }\n  }\n':
     types.RegisterDocument,
+  '\n  query GetIsAdminInitialized {\n    isAdminInitialized\n  }\n':
+    types.GetIsAdminInitializedDocument,
   '\n  query GetWorkers {\n    workers {\n      kind\n      name\n      addr\n      device\n      arch\n      cpuInfo\n      cpuCount\n      cudaDevices\n    }\n  }\n':
     types.GetWorkersDocument,
   '\n  mutation refreshToken($refreshToken: String!) {\n    refreshToken(refreshToken: $refreshToken) {\n      accessToken\n      refreshToken\n    }\n  }\n':
@@ -49,14 +51,20 @@ export function graphql(
  * The graphql function is used to parse GraphQL queries into a document that can be used by GraphQL clients.
  */
 export function graphql(
-  source: '\n  query GetIsAdminInitialized {\n    isAdminInitialized\n  }\n'
-): (typeof documents)['\n  query GetIsAdminInitialized {\n    isAdminInitialized\n  }\n']
+  source: '\n  mutation tokenAuth($email: String!, $password: String!) {\n    tokenAuth(email: $email, password: $password) {\n      accessToken\n      refreshToken\n    }\n  }\n'
+): (typeof documents)['\n  mutation tokenAuth($email: String!, $password: String!) {\n    tokenAuth(email: $email, password: $password) {\n      accessToken\n      refreshToken\n    }\n  }\n']
 /**
  * The graphql function is used to parse GraphQL queries into a document that can be used by GraphQL clients.
  */
 export function graphql(
   source: '\n  mutation register(\n    $email: String!\n    $password1: String!\n    $password2: String!\n    $invitationCode: String\n  ) {\n    register(\n      email: $email\n      password1: $password1\n      password2: $password2\n      invitationCode: $invitationCode\n    ) {\n      accessToken\n      refreshToken\n    }\n  }\n'
 ): (typeof documents)['\n  mutation register(\n    $email: String!\n    $password1: String!\n    $password2: String!\n    $invitationCode: String\n  ) {\n    register(\n      email: $email\n      password1: $password1\n      password2: $password2\n      invitationCode: $invitationCode\n    ) {\n      accessToken\n      refreshToken\n    }\n  }\n']
+/**
+ * The graphql function is used to parse GraphQL queries into a document that can be used by GraphQL clients.
+ */
+export function graphql(
+  source: '\n  query GetIsAdminInitialized {\n    isAdminInitialized\n  }\n'
+): (typeof documents)['\n  query GetIsAdminInitialized {\n    isAdminInitialized\n  }\n']
 /**
  * The graphql function is used to parse GraphQL queries into a document that can be used by GraphQL clients.
  */
diff --git a/ee/tabby-ui/lib/gql/generates/graphql.ts b/ee/tabby-ui/lib/gql/generates/graphql.ts
index 17e5282..1d821bd 100644
--- a/ee/tabby-ui/lib/gql/generates/graphql.ts
+++ b/ee/tabby-ui/lib/gql/generates/graphql.ts
@@ -146,13 +146,18 @@ export type GetRegistrationTokenQuery = {
   registrationToken: string
 }
 
-export type GetIsAdminInitializedQueryVariables = Exact<{
-  [key: string]: never
+export type TokenAuthMutationVariables = Exact<{
+  email: Scalars['String']['input']
+  password: Scalars['String']['input']
 }>
 
-export type GetIsAdminInitializedQuery = {
-  __typename?: 'Query'
-  isAdminInitialized: boolean
+export type TokenAuthMutation = {
+  __typename?: 'Mutation'
+  tokenAuth: {
+    __typename?: 'TokenAuthResponse'
+    accessToken: string
+    refreshToken: string
+  }
 }
 
 export type RegisterMutationVariables = Exact<{
@@ -171,6 +176,15 @@ export type RegisterMutation = {
   }
 }
 
+export type GetIsAdminInitializedQueryVariables = Exact<{
+  [key: string]: never
+}>
+
+export type GetIsAdminInitializedQuery = {
+  __typename?: 'Query'
+  isAdminInitialized: boolean
+}
+
 export type GetWorkersQueryVariables = Exact<{ [key: string]: never }>
 
 export type GetWorkersQuery = {
@@ -220,25 +234,74 @@ export const GetRegistrationTokenDocument = {
   GetRegistrationTokenQuery,
   GetRegistrationTokenQueryVariables
 >
-export const GetIsAdminInitializedDocument = {
+export const TokenAuthDocument = {
   kind: 'Document',
   definitions: [
     {
       kind: 'OperationDefinition',
-      operation: 'query',
-      name: { kind: 'Name', value: 'GetIsAdminInitialized' },
+      operation: 'mutation',
+      name: { kind: 'Name', value: 'tokenAuth' },
+      variableDefinitions: [
+        {
+          kind: 'VariableDefinition',
+          variable: {
+            kind: 'Variable',
+            name: { kind: 'Name', value: 'email' }
+          },
+          type: {
+            kind: 'NonNullType',
+            type: { kind: 'NamedType', name: { kind: 'Name', value: 'String' } }
+          }
+        },
+        {
+          kind: 'VariableDefinition',
+          variable: {
+            kind: 'Variable',
+            name: { kind: 'Name', value: 'password' }
+          },
+          type: {
+            kind: 'NonNullType',
+            type: { kind: 'NamedType', name: { kind: 'Name', value: 'String' } }
+          }
+        }
+      ],
       selectionSet: {
         kind: 'SelectionSet',
         selections: [
-          { kind: 'Field', name: { kind: 'Name', value: 'isAdminInitialized' } }
+          {
+            kind: 'Field',
+            name: { kind: 'Name', value: 'tokenAuth' },
+            arguments: [
+              {
+                kind: 'Argument',
+                name: { kind: 'Name', value: 'email' },
+                value: {
+                  kind: 'Variable',
+                  name: { kind: 'Name', value: 'email' }
+                }
+              },
+              {
+                kind: 'Argument',
+                name: { kind: 'Name', value: 'password' },
+                value: {
+                  kind: 'Variable',
+                  name: { kind: 'Name', value: 'password' }
+                }
+              }
+            ],
+            selectionSet: {
+              kind: 'SelectionSet',
+              selections: [
+                { kind: 'Field', name: { kind: 'Name', value: 'accessToken' } },
+                { kind: 'Field', name: { kind: 'Name', value: 'refreshToken' } }
+              ]
+            }
+          }
         ]
       }
     }
   ]
-} as unknown as DocumentNode<
-  GetIsAdminInitializedQuery,
-  GetIsAdminInitializedQueryVariables
->
+} as unknown as DocumentNode<TokenAuthMutation, TokenAuthMutationVariables>
 export const RegisterDocument = {
   kind: 'Document',
   definitions: [
@@ -342,6 +405,25 @@ export const RegisterDocument = {
     }
   ]
 } as unknown as DocumentNode<RegisterMutation, RegisterMutationVariables>
+export const GetIsAdminInitializedDocument = {
+  kind: 'Document',
+  definitions: [
+    {
+      kind: 'OperationDefinition',
+      operation: 'query',
+      name: { kind: 'Name', value: 'GetIsAdminInitialized' },
+      selectionSet: {
+        kind: 'SelectionSet',
+        selections: [
+          { kind: 'Field', name: { kind: 'Name', value: 'isAdminInitialized' } }
+        ]
+      }
+    }
+  ]
+} as unknown as DocumentNode<
+  GetIsAdminInitializedQuery,
+  GetIsAdminInitializedQueryVariables
+>
 export const GetWorkersDocument = {
   kind: 'Document',
   definitions: [
diff --git a/ee/tabby-ui/lib/tabby/auth.tsx b/ee/tabby-ui/lib/tabby/auth.tsx
index 6087808..5d55c67 100644
--- a/ee/tabby-ui/lib/tabby/auth.tsx
+++ b/ee/tabby-ui/lib/tabby/auth.tsx
@@ -142,7 +142,11 @@ const AuthProvider: React.FunctionComponent<AuthProviderProps> = ({
     data: null
   })
 
+  const initialized = React.useRef(false)
   React.useEffect(() => {
+    if (initialized.current) return
+
+    initialized.current = true
     const data = storage.initialState()
     if (data?.refreshToken) {
       doRefresh(data.refreshToken, dispatch)
@@ -219,9 +223,14 @@ type Session =
 function useSession(): Session {
   const { authState } = useAuthStore()
   if (authState?.status == 'authenticated') {
-    const { user } = jwtDecode<{ user: User }>(authState.data.accessToken)
+    const { user } = jwtDecode<{ user: { email: string; is_admin: boolean } }>(
+      authState.data.accessToken
+    )
     return {
-      data: user,
+      data: {
+        email: user.email,
+        isAdmin: user.is_admin
+      },
       status: authState.status
     }
   } else {
diff --git a/ee/tabby-ui/lib/tabby/gql.ts b/ee/tabby-ui/lib/tabby/gql.ts
index 316dce1..1e6aedc 100644
--- a/ee/tabby-ui/lib/tabby/gql.ts
+++ b/ee/tabby-ui/lib/tabby/gql.ts
@@ -22,12 +22,14 @@ export function useGraphQLForm<
 >(
   document: TypedDocumentNode<TResult, TVariables>,
   options?: {
+    onSuccess?: (values: TResult) => void
     onError?: (path: string, message: string) => void
   }
 ) {
   const onSubmit = async (values: TVariables) => {
+    let res
     try {
-      await gqlClient.request(document, values)
+      res = await gqlClient.request(document, values)
     } catch (err) {
       const { errors = [] } = (err as any).response as GraphQLResponse
       for (const error of errors) {
@@ -42,7 +44,11 @@ export function useGraphQLForm<
           options?.onError && options?.onError('root', error.message)
         }
       }
+
+      return res
     }
+
+    options?.onSuccess && options.onSuccess(res)
   }
   return { onSubmit }
 }