Merge pull request #4146 from janhq/fix/app-rerender-issues

fix: app re-render issues caused by bad state handling

web/containers/Layout/index.tsx

@@ -1,10 +1,8 @@
 'use client'
 
-import { useEffect } from 'react'
+import { useEffect, useMemo } from 'react'
 
-import { motion as m } from 'framer-motion'
-
-import { useAtom, useAtomValue } from 'jotai'
+import { useAtomValue, useSetAtom } from 'jotai'
 
 import { twMerge } from 'tailwind-merge'
 
@@ -36,7 +34,7 @@ import { mainViewStateAtom } from '@/helpers/atoms/App.atom'
 import { reduceTransparentAtom } from '@/helpers/atoms/Setting.atom'
 
 const BaseLayout = () => {
-  const [mainViewState, setMainViewState] = useAtom(mainViewStateAtom)
+  const setMainViewState = useSetAtom(mainViewStateAtom)
   const importModelStage = useAtomValue(getImportModelStageAtom)
   const reduceTransparent = useAtomValue(reduceTransparentAtom)
 
@@ -68,24 +66,7 @@ const BaseLayout = () => {
       <TopPanel />
       <div className="relative top-9 flex h-[calc(100vh-(36px+36px))] w-screen">
         <RibbonPanel />
-        <div className={twMerge('relative flex w-full')}>
-          <div className="w-full">
-            <m.div
-              key={mainViewState}
-              initial={{ opacity: 0, y: -8 }}
-              className="h-full"
-              animate={{
-                opacity: 1,
-                y: 0,
-                transition: {
-                  duration: 0.5,
-                },
-              }}
-            >
-              <MainViewContainer />
-            </m.div>
-          </div>
-        </div>
+        <MainViewContainer />
         <LoadingModal />
         {importModelStage === 'SELECTING_MODEL' && <SelectingModelModal />}
         {importModelStage === 'MODEL_SELECTED' && <ImportModelOptionModal />}

web/containers/MainViewContainer/index.tsx

@@ -1,5 +1,10 @@
+import { memo } from 'react'
+
+import { motion as m } from 'framer-motion'
 import { useAtomValue } from 'jotai'
 
+import { twMerge } from 'tailwind-merge'
+
 import { MainViewState } from '@/constants/screens'
 
 import HubScreen from '@/screens/Hub'
@@ -31,7 +36,26 @@ const MainViewContainer = () => {
       break
   }
 
-  return children
+  return (
+    <div className={twMerge('relative flex w-full')}>
+      <div className="w-full">
+        <m.div
+          key={mainViewState}
+          initial={{ opacity: 0, y: -8 }}
+          className="h-full"
+          animate={{
+            opacity: 1,
+            y: 0,
+            transition: {
+              duration: 0.25,
+            },
+          }}
+        >
+          {children}
+        </m.div>
+      </div>
+    </div>
+  )
 }
 
-export default MainViewContainer
+export default memo(MainViewContainer)

web/containers/Providers/CoreConfigurator.tsx

@@ -0,0 +1,64 @@
+'use client'
+
+import { PropsWithChildren, useCallback, useEffect, useState } from 'react'
+
+import Loader from '@/containers/Loader'
+
+import { setupCoreServices } from '@/services/coreService'
+import {
+  isCoreExtensionInstalled,
+  setupBaseExtensions,
+} from '@/services/extensionService'
+
+import { extensionManager } from '@/extension'
+
+export const CoreConfigurator = ({ children }: PropsWithChildren) => {
+  const [setupCore, setSetupCore] = useState(false)
+  const [activated, setActivated] = useState(false)
+  const [settingUp, setSettingUp] = useState(false)
+
+  const setupExtensions = useCallback(async () => {
+    // Register all active extensions
+    await extensionManager.registerActive()
+
+    setTimeout(async () => {
+      if (!isCoreExtensionInstalled()) {
+        setSettingUp(true)
+        await setupBaseExtensions()
+        return
+      }
+
+      extensionManager.load()
+      setSettingUp(false)
+      setActivated(true)
+    }, 500)
+  }, [])
+
+  // Services Setup
+  useEffect(() => {
+    setupCoreServices()
+    setSetupCore(true)
+    return () => {
+      extensionManager.unload()
+    }
+  }, [])
+
+  useEffect(() => {
+    if (setupCore) {
+      // Electron
+      if (window && window.core?.api) {
+        setupExtensions()
+      } else {
+        // Host
+        setActivated(true)
+      }
+    }
+  }, [setupCore, setupExtensions])
+
+  return (
+    <>
+      {settingUp && <Loader description="Preparing Update..." />}
+      {setupCore && activated && <>{children}</>}
+    </>
+  )
+}

web/containers/Providers/index.tsx

@@ -1,23 +1,17 @@
 'use client'
 
-import { PropsWithChildren, useCallback, useEffect, useState } from 'react'
+import { PropsWithChildren } from 'react'
 
 import { Toaster } from 'react-hot-toast'
 
-import Loader from '@/containers/Loader'
 import EventListener from '@/containers/Providers/EventListener'
 import JotaiWrapper from '@/containers/Providers/Jotai'
 
 import ThemeWrapper from '@/containers/Providers/Theme'
 
-import { setupCoreServices } from '@/services/coreService'
-import {
-  isCoreExtensionInstalled,
-  setupBaseExtensions,
-} from '@/services/extensionService'
-
 import Umami from '@/utils/umami'
 
+import { CoreConfigurator } from './CoreConfigurator'
 import DataLoader from './DataLoader'
 
 import DeepLinkListener from './DeepLinkListener'
@@ -26,57 +20,12 @@ import Responsive from './Responsive'
 
 import SettingsHandler from './SettingsHandler'
 
-import { extensionManager } from '@/extension'
-
 const Providers = ({ children }: PropsWithChildren) => {
-  const [setupCore, setSetupCore] = useState(false)
-  const [activated, setActivated] = useState(false)
-  const [settingUp, setSettingUp] = useState(false)
-
-  const setupExtensions = useCallback(async () => {
-    // Register all active extensions
-    await extensionManager.registerActive()
-
-    setTimeout(async () => {
-      if (!isCoreExtensionInstalled()) {
-        setSettingUp(true)
-        await setupBaseExtensions()
-        return
-      }
-
-      extensionManager.load()
-      setSettingUp(false)
-      setActivated(true)
-    }, 500)
-  }, [])
-
-  // Services Setup
-  useEffect(() => {
-    setupCoreServices()
-    setSetupCore(true)
-    return () => {
-      extensionManager.unload()
-    }
-  }, [])
-
-  useEffect(() => {
-    if (setupCore) {
-      // Electron
-      if (window && window.core?.api) {
-        setupExtensions()
-      } else {
-        // Host
-        setActivated(true)
-      }
-    }
-  }, [setupCore, setupExtensions])
-
   return (
     <ThemeWrapper>
       <JotaiWrapper>
         <Umami />
-        {settingUp && <Loader description="Preparing Update..." />}
-        {setupCore && activated && (
+        <CoreConfigurator>
           <>
             <Responsive />
             <KeyListener />
@@ -87,7 +36,7 @@ const Providers = ({ children }: PropsWithChildren) => {
             <Toaster />
             {children}
           </>
-        )}
+        </CoreConfigurator>
       </JotaiWrapper>
     </ThemeWrapper>
   )

web/helpers/atoms/AppConfig.atom.ts

@@ -12,14 +12,35 @@ export const janDataFolderPathAtom = atom('')
 
 export const experimentalFeatureEnabledAtom = atomWithStorage(
   EXPERIMENTAL_FEATURE,
-  false
+  false,
+  undefined,
+  { getOnInit: true }
 )
 
-export const proxyEnabledAtom = atomWithStorage(PROXY_FEATURE_ENABLED, false)
-export const proxyAtom = atomWithStorage(HTTPS_PROXY_FEATURE, '')
+export const proxyEnabledAtom = atomWithStorage(
+  PROXY_FEATURE_ENABLED,
+  false,
+  undefined,
+  { getOnInit: true }
+)
+export const proxyAtom = atomWithStorage(HTTPS_PROXY_FEATURE, '', undefined, {
+  getOnInit: true,
+})
 
-export const ignoreSslAtom = atomWithStorage(IGNORE_SSL, false)
-export const vulkanEnabledAtom = atomWithStorage(VULKAN_ENABLED, false)
-export const quickAskEnabledAtom = atomWithStorage(QUICK_ASK_ENABLED, false)
+export const ignoreSslAtom = atomWithStorage(IGNORE_SSL, false, undefined, {
+  getOnInit: true,
+})
+export const vulkanEnabledAtom = atomWithStorage(
+  VULKAN_ENABLED,
+  false,
+  undefined,
+  { getOnInit: true }
+)
+export const quickAskEnabledAtom = atomWithStorage(
+  QUICK_ASK_ENABLED,
+  false,
+  undefined,
+  { getOnInit: true }
+)
 
 export const hostAtom = atom('http://localhost:1337/')

web/helpers/atoms/Model.atom.ts

@@ -16,7 +16,9 @@ enum ModelStorageAtomKeys {
  */
 export const downloadedModelsAtom = atomWithStorage<Model[]>(
   ModelStorageAtomKeys.DownloadedModels,
-  []
+  [],
+  undefined,
+  { getOnInit: true }
 )
 
 /**
@@ -25,7 +27,9 @@ export const downloadedModelsAtom = atomWithStorage<Model[]>(
  */
 export const configuredModelsAtom = atomWithStorage<Model[]>(
   ModelStorageAtomKeys.AvailableModels,
-  []
+  [],
+  undefined,
+  { getOnInit: true }
 )
 
 export const removeDownloadedModelAtom = atom(

web/helpers/atoms/Setting.atom.ts

@@ -13,10 +13,22 @@ export const REDUCE_TRANSPARENT = 'reduceTransparent'
 export const SPELL_CHECKING = 'spellChecking'
 export const themesOptionsAtom = atom<{ name: string; value: string }[]>([])
 export const janThemesPathAtom = atom<string | undefined>(undefined)
-export const selectedThemeIdAtom = atomWithStorage<string>(THEME, '')
+export const selectedThemeIdAtom = atomWithStorage<string>(
+  THEME,
+  '',
+  undefined,
+  { getOnInit: true }
+)
 export const themeDataAtom = atom<Theme | undefined>(undefined)
 export const reduceTransparentAtom = atomWithStorage<boolean>(
   REDUCE_TRANSPARENT,
-  false
+  false,
+  undefined,
+  { getOnInit: true }
+)
+export const spellCheckAtom = atomWithStorage<boolean>(
+  SPELL_CHECKING,
+  false,
+  undefined,
+  { getOnInit: true }
 )
-export const spellCheckAtom = atomWithStorage<boolean>(SPELL_CHECKING, false)

web/helpers/atoms/Thread.atom.ts

@@ -207,7 +207,9 @@ export const setThreadModelParamsAtom = atom(
  */
 export const activeSettingInputBoxAtom = atomWithStorage<boolean>(
   ACTIVE_SETTING_INPUT_BOX,
-  false
+  false,
+  undefined,
+  { getOnInit: true }
 )
 
 /**

web/hooks/useStarterScreen.ts

@@ -1,22 +1,20 @@
-import { useState, useEffect } from 'react'
+import { useState, useEffect, useMemo } from 'react'
 
-import { useAtomValue, useSetAtom } from 'jotai'
+import { useAtomValue } from 'jotai'
 
 import { isLocalEngine } from '@/utils/modelEngine'
 
 import { extensionManager } from '@/extension'
-import {
-  downloadedModelsAtom,
-  selectedModelAtom,
-} from '@/helpers/atoms/Model.atom'
+import { downloadedModelsAtom } from '@/helpers/atoms/Model.atom'
 import { threadsAtom } from '@/helpers/atoms/Thread.atom'
 
 export function useStarterScreen() {
   const downloadedModels = useAtomValue(downloadedModelsAtom)
   const threads = useAtomValue(threadsAtom)
-  const setSelectedModel = useSetAtom(selectedModelAtom)
-  const isDownloadALocalModel = downloadedModels.some((x) =>
-    isLocalEngine(x.engine)
+
+  const isDownloadALocalModel = useMemo(
+    () => downloadedModels.some((x) => isLocalEngine(x.engine)),
+    [downloadedModels]
   )
 
   const [extensionHasSettings, setExtensionHasSettings] = useState<
@@ -24,9 +22,6 @@ export function useStarterScreen() {
   >([])
 
   useEffect(() => {
-    if (isDownloadALocalModel) {
-      setSelectedModel(downloadedModels[0])
-    }
     const getAllSettings = async () => {
       const extensionsMenu: {
         name?: string
@@ -66,12 +61,16 @@ export function useStarterScreen() {
     // eslint-disable-next-line react-hooks/exhaustive-deps
   }, [])
 
-  const isAnyRemoteModelConfigured = extensionHasSettings.some(
-    (x) => x.apiKey.length > 1
+  const isAnyRemoteModelConfigured = useMemo(
+    () => extensionHasSettings.some((x) => x.apiKey.length > 1),
+    [extensionHasSettings]
   )
 
-  const isShowStarterScreen =
-    !isAnyRemoteModelConfigured && !isDownloadALocalModel && !threads.length
+  const isShowStarterScreen = useMemo(
+    () =>
+      !isAnyRemoteModelConfigured && !isDownloadALocalModel && !threads.length,
+    [isAnyRemoteModelConfigured, isDownloadALocalModel, threads]
+  )
 
   return {
     extensionHasSettings,

web/screens/Thread/ThreadCenterPanel/ChatBody/OnDeviceStarterScreen/index.tsx

@@ -24,6 +24,8 @@ import useDownloadModel from '@/hooks/useDownloadModel'
 
 import { modelDownloadStateAtom } from '@/hooks/useDownloadState'
 
+import { useStarterScreen } from '@/hooks/useStarterScreen'
+
 import { formatDownloadPercentage, toGibibytes } from '@/utils/converter'
 import {
   getLogoEngine,
@@ -38,16 +40,8 @@ import {
 } from '@/helpers/atoms/Model.atom'
 import { selectedSettingAtom } from '@/helpers/atoms/Setting.atom'
 
-type Props = {
-  extensionHasSettings: {
-    name?: string
-    setting: string
-    apiKey: string
-    provider: string
-  }[]
-}
-
-const OnDeviceStarterScreen = ({ extensionHasSettings }: Props) => {
+const OnDeviceStarterScreen = () => {
+  const { extensionHasSettings } = useStarterScreen()
   const [searchValue, setSearchValue] = useState('')
   const [isOpen, setIsOpen] = useState(Boolean(searchValue.length))
   const downloadingModels = useAtomValue(getDownloadingModelAtom)

web/screens/Thread/ThreadCenterPanel/ChatBody/index.tsx

@@ -1,3 +1,5 @@
+import { memo } from 'react'
+
 import { MessageStatus } from '@janhq/core'
 
 import { useAtomValue } from 'jotai'
@@ -44,4 +46,4 @@ const ChatBody = () => {
   )
 }
 
-export default ChatBody
+export default memo(ChatBody)

web/screens/Thread/ThreadCenterPanel/index.tsx

@@ -1,6 +1,6 @@
 /* eslint-disable @typescript-eslint/naming-convention */
 
-import { useEffect, useState } from 'react'
+import { memo, useEffect, useState } from 'react'
 
 import { Accept, useDropzone } from 'react-dropzone'
 
@@ -232,4 +232,4 @@ const ThreadCenterPanel = () => {
   )
 }
 
-export default ThreadCenterPanel
+export default memo(ThreadCenterPanel)

web/screens/Thread/index.tsx

@@ -1,3 +1,5 @@
+import { memo } from 'react'
+
 import { useStarterScreen } from '@/hooks/useStarterScreen'
 
 import ThreadLeftPanel from '@/screens/Thread/ThreadLeftPanel'
@@ -9,19 +11,31 @@ import ModalDeleteThread from './ThreadLeftPanel/ModalDeleteThread'
 import ModalEditTitleThread from './ThreadLeftPanel/ModalEditTitleThread'
 import ThreadRightPanel from './ThreadRightPanel'
 
+type Props = {
+  isShowStarterScreen: boolean
+}
+
+const ThreadPanels = memo(({ isShowStarterScreen }: Props) => {
+  return isShowStarterScreen ? (
+    <OnDeviceStarterScreen />
+  ) : (
+    <>
+      <ThreadLeftPanel />
+      <ThreadCenterPanel />
+      <ThreadRightPanel />
+    </>
+  )
+})
+
+const WelcomeController = () => {
+  const { isShowStarterScreen } = useStarterScreen()
+  return <ThreadPanels isShowStarterScreen={isShowStarterScreen} />
+}
+
 const ThreadScreen = () => {
-  const { extensionHasSettings, isShowStarterScreen } = useStarterScreen()
   return (
     <div className="relative flex h-full w-full flex-1 overflow-x-hidden">
-      {isShowStarterScreen ? (
-        <OnDeviceStarterScreen extensionHasSettings={extensionHasSettings} />
-      ) : (
-        <>
-          <ThreadLeftPanel />
-          <ThreadCenterPanel />
-          <ThreadRightPanel />
-        </>
-      )}
+      <WelcomeController />
 
       {/* Showing variant modal action for thread screen */}
       <ModalEditTitleThread />
@@ -31,4 +45,4 @@ const ThreadScreen = () => {
   )
 }
 
-export default ThreadScreen
+export default memo(ThreadScreen)