From e621a5a418798c59c3f2a9559cc2d483601d44a1 Mon Sep 17 00:00:00 2001
From: Meng Zhang <meng@tabbyml.com>
Date: Sun, 18 Feb 2024 00:13:31 -0800
Subject: [PATCH] fix: use 'serverInfo' to check whether chat is enabled, avoid
 using admin access query 'workers' (#1477)

---
 ee/tabby-ui/components/user-panel.tsx    |  7 ++----
 ee/tabby-ui/lib/hooks/use-server-info.ts | 32 ++++++++++++++++++++++++
 ee/tabby-ui/lib/tabby/auth.tsx           | 16 +++++-------
 3 files changed, 40 insertions(+), 15 deletions(-)
 create mode 100644 ee/tabby-ui/lib/hooks/use-server-info.ts

diff --git a/ee/tabby-ui/components/user-panel.tsx b/ee/tabby-ui/components/user-panel.tsx
index 239a22e67515..693109abe951 100644
--- a/ee/tabby-ui/components/user-panel.tsx
+++ b/ee/tabby-ui/components/user-panel.tsx
@@ -1,9 +1,7 @@
 import React from 'react'
-import { has } from 'lodash-es'
 import NiceAvatar, { genConfig } from 'react-nice-avatar'
 
-import { WorkerKind } from '@/lib/gql/generates/graphql'
-import { useWorkers } from '@/lib/hooks/use-workers'
+import { useIsChatEnabled } from '@/lib/hooks/use-server-info'
 import { useAuthenticatedSession, useSignOut } from '@/lib/tabby/auth'
 import {
   DropdownMenu,
@@ -20,8 +18,7 @@ export default function UserPanel() {
   const user = useAuthenticatedSession()
   const signOut = useSignOut()
 
-  const workers = useWorkers()
-  const isChatEnabled = has(workers, WorkerKind.Chat)
+  const isChatEnabled = useIsChatEnabled()
 
   if (!user) {
     return
diff --git a/ee/tabby-ui/lib/hooks/use-server-info.ts b/ee/tabby-ui/lib/hooks/use-server-info.ts
new file mode 100644
index 000000000000..45ba6fc7a1e1
--- /dev/null
+++ b/ee/tabby-ui/lib/hooks/use-server-info.ts
@@ -0,0 +1,32 @@
+import { useQuery } from 'urql'
+
+import { graphql } from '@/lib/gql/generates'
+
+export const getServerInfo = graphql(/* GraphQL */ `
+  query GetServerInfo {
+    serverInfo {
+      isAdminInitialized
+      isEmailConfigured
+      isChatEnabled
+    }
+  }
+`)
+
+const useServerInfo = () => {
+  const [{ data }] = useQuery({ query: getServerInfo })
+  return data?.serverInfo
+}
+
+const useIsChatEnabled = () => {
+  return useServerInfo()?.isChatEnabled
+}
+
+const useIsAdminInitialized = () => {
+  return useServerInfo()?.isAdminInitialized
+}
+
+const useIsEmailConfigured = () => {
+  return useServerInfo()?.isEmailConfigured
+}
+
+export { useIsChatEnabled, useIsAdminInitialized, useIsEmailConfigured }
diff --git a/ee/tabby-ui/lib/tabby/auth.tsx b/ee/tabby-ui/lib/tabby/auth.tsx
index d65fe142e28c..40e134b43668 100644
--- a/ee/tabby-ui/lib/tabby/auth.tsx
+++ b/ee/tabby-ui/lib/tabby/auth.tsx
@@ -1,12 +1,13 @@
 import * as React from 'react'
 import { usePathname, useRouter } from 'next/navigation'
 import { jwtDecode, JwtPayload } from 'jwt-decode'
-import { useQuery } from 'urql'
 import useLocalStorage from 'use-local-storage'
 
 import { graphql } from '@/lib/gql/generates'
 import { isClientSide } from '@/lib/utils'
 
+import { useIsAdminInitialized } from '../hooks/use-server-info'
+
 interface AuthData {
   accessToken: string
   refreshToken: string
@@ -279,16 +280,10 @@ function useSession(): Session {
   return session
 }
 
-export const getIsAdminInitialized = graphql(/* GraphQL */ `
-  query GetIsAdminInitialized {
-    isAdminInitialized
-  }
-`)
-
 const redirectWhitelist = ['/auth/signin', '/auth/signup']
 
 function useAuthenticatedSession() {
-  const [{ data }] = useQuery({ query: getIsAdminInitialized })
+  const isAdminInitialized = useIsAdminInitialized()
   const router = useRouter()
   const pathName = usePathname()
   const { data: session, status } = useSession()
@@ -296,13 +291,14 @@ function useAuthenticatedSession() {
   React.useEffect(() => {
     if (status === 'loading') return
     if (status === 'authenticated') return
+    if (isAdminInitialized === undefined) return
 
-    if (data?.isAdminInitialized === false) {
+    if (!isAdminInitialized) {
       router.replace('/auth/signup?isAdmin=true')
     } else if (!redirectWhitelist.includes(pathName)) {
       router.replace('/auth/signin')
     }
-  }, [data, status])
+  }, [isAdminInitialized, status])
 
   return session
 }