Improve Jotai work (#18205)

packages/twenty-front/src/modules/app/components/App.tsx

@@ -6,9 +6,9 @@ import { AppRootErrorFallback } from '@/error-handler/components/AppRootErrorFal
 import { ExceptionHandlerProvider } from '@/error-handler/components/ExceptionHandlerProvider';
 import { SnackBarComponentInstanceContext } from '@/ui/feedback/snack-bar-manager/contexts/SnackBarComponentInstanceContext';
 import { ClickOutsideListenerContext } from '@/ui/utilities/pointer-event/contexts/ClickOutsideListenerContext';
-import { jotaiStore } from '@/ui/utilities/state/jotai/jotaiStore';
 import { i18n } from '@lingui/core';
 import { I18nProvider } from '@lingui/react';
+import { jotaiStore } from '@/ui/utilities/state/jotai/jotaiStore';
 import { Provider as JotaiProvider } from 'jotai';
 import { HelmetProvider } from 'react-helmet-async';
 import { RecoilRoot } from 'recoil';

packages/twenty-front/src/modules/app/effect-components/PageChangeEffect.tsx

@@ -35,7 +35,6 @@ import { PageFocusId } from '@/types/PageFocusId';
 import { useResetFocusStackToFocusItem } from '@/ui/utilities/focus/hooks/useResetFocusStackToFocusItem';
 import { FocusComponentType } from '@/ui/utilities/focus/types/FocusComponentType';
 import { useRecoilComponentValueV2 } from '@/ui/utilities/state/jotai/hooks/useRecoilComponentValueV2';
-import { jotaiStore } from '@/ui/utilities/state/jotai/jotaiStore';
 import { AppBasePath, AppPath, CommandMenuPages } from 'twenty-shared/types';
 import { isDefined } from 'twenty-shared/utils';
 import { AnalyticsType } from '~/generated-metadata/graphql';
@@ -43,10 +42,12 @@ import { usePageChangeEffectNavigateLocation } from '~/hooks/usePageChangeEffect
 import { useInitializeQueryParamState } from '~/modules/app/hooks/useInitializeQueryParamState';
 import { isMatchingLocation } from '~/utils/isMatchingLocation';
 import { getPageTitleFromPath } from '~/utils/title-utils';
+import { useStore } from 'jotai';
 
 // TODO: break down into smaller functions and / or hooks
 //  - moved usePageChangeEffectNavigateLocation into dedicated hook
 export const PageChangeEffect = () => {
+  const store = useStore();
   const navigate = useNavigate();
 
   const [previousLocation, setPreviousLocation] = useState('');
@@ -99,12 +100,12 @@ export const PageChangeEffect = () => {
   const { closeCommandMenu } = useCommandMenu();
 
   const closeCommandMenuUnlessOnEditPage = useCallback(() => {
-    const currentPage = jotaiStore.get(commandMenuPageState.atom);
+    const currentPage = store.get(commandMenuPageState.atom);
     if (currentPage === CommandMenuPages.NavigationMenuItemEdit) {
       return;
     }
     closeCommandMenu();
-  }, [closeCommandMenu]);
+  }, [closeCommandMenu, store]);
 
   const { resetFocusStackToFocusItem } = useResetFocusStackToFocusItem();
 

packages/twenty-front/src/modules/app/hooks/useExecuteTasksOnAnyLocationChange.ts

@@ -1,7 +1,6 @@
 import { MAIN_CONTEXT_STORE_INSTANCE_ID } from '@/context-store/constants/MainContextStoreInstanceId';
 import { contextStoreIsPageInEditModeComponentState } from '@/context-store/states/contextStoreIsPageInEditModeComponentState';
 import { currentPageLayoutIdState } from '@/page-layout/states/currentPageLayoutIdState';
-import { jotaiStore } from '@/ui/utilities/state/jotai/jotaiStore';
 import { hasInitializedFieldsWidgetGroupsDraftComponentState } from '@/page-layout/states/hasInitializedFieldsWidgetGroupsDraftComponentState';
 import { isPageLayoutInEditModeComponentState } from '@/page-layout/states/isPageLayoutInEditModeComponentState';
 import { pageLayoutCurrentLayoutsComponentState } from '@/page-layout/states/pageLayoutCurrentLayoutsComponentState';
@@ -12,22 +11,24 @@ import { convertPageLayoutToTabLayouts } from '@/page-layout/utils/convertPageLa
 import { useCloseAnyOpenDropdown } from '@/ui/layout/dropdown/hooks/useCloseAnyOpenDropdown';
 import { useCallback } from 'react';
 import { isDefined } from 'twenty-shared/utils';
+import { useStore } from 'jotai';
 
 export const useExecuteTasksOnAnyLocationChange = () => {
+  const store = useStore();
   const { closeAnyOpenDropdown } = useCloseAnyOpenDropdown();
 
   const resetPageLayoutEditMode = useCallback(() => {
-    const pageLayoutId = jotaiStore.get(currentPageLayoutIdState.atom);
+    const pageLayoutId = store.get(currentPageLayoutIdState.atom);
 
     if (isDefined(pageLayoutId)) {
-      const pageLayoutPersisted = jotaiStore.get(
+      const pageLayoutPersisted = store.get(
         pageLayoutPersistedComponentState.atomFamily({
           instanceId: pageLayoutId,
         }),
       );
 
       if (isDefined(pageLayoutPersisted)) {
-        jotaiStore.set(
+        store.set(
           pageLayoutDraftComponentState.atomFamily({
             instanceId: pageLayoutId,
           }),
@@ -41,7 +42,7 @@ export const useExecuteTasksOnAnyLocationChange = () => {
         );
 
         const tabLayouts = convertPageLayoutToTabLayouts(pageLayoutPersisted);
-        jotaiStore.set(
+        store.set(
           pageLayoutCurrentLayoutsComponentState.atomFamily({
             instanceId: pageLayoutId,
           }),
@@ -49,37 +50,37 @@ export const useExecuteTasksOnAnyLocationChange = () => {
         );
       }
 
-      jotaiStore.set(
+      store.set(
         isPageLayoutInEditModeComponentState.atomFamily({
           instanceId: pageLayoutId,
         }),
         false,
       );
 
-      jotaiStore.set(
+      store.set(
         pageLayoutIsInitializedComponentState.atomFamily({
           instanceId: pageLayoutId,
         }),
         false,
       );
 
-      jotaiStore.set(
+      store.set(
         hasInitializedFieldsWidgetGroupsDraftComponentState.atomFamily({
           instanceId: pageLayoutId,
         }),
         {},
       );
 
-      jotaiStore.set(currentPageLayoutIdState.atom, null);
+      store.set(currentPageLayoutIdState.atom, null);
     }
 
-    jotaiStore.set(
+    store.set(
       contextStoreIsPageInEditModeComponentState.atomFamily({
         instanceId: MAIN_CONTEXT_STORE_INSTANCE_ID,
       }),
       false,
     );
-  }, []);
+  }, [store]);
 
   /**
    * Be careful to put idempotent tasks here.

packages/twenty-front/src/modules/app/hooks/useInitializeQueryParamState.ts

@@ -3,16 +3,17 @@ import { useCallback } from 'react';
 import { billingCheckoutSessionState } from '@/auth/states/billingCheckoutSessionState';
 import { type BillingCheckoutSession } from '@/auth/types/billingCheckoutSession.type';
 import { BILLING_CHECKOUT_SESSION_DEFAULT_VALUE } from '@/billing/constants/BillingCheckoutSessionDefaultValue';
-import { jotaiStore } from '@/ui/utilities/state/jotai/jotaiStore';
 import deepEqual from 'deep-equal';
+import { useStore } from 'jotai';
 
 // Initialize state that are hydrated from query parameters
 // We used to use recoil-sync to do this, but it was causing issues with Firefox
 export const useInitializeQueryParamState = () => {
+  const store = useStore();
   const initializeQueryParamState = useCallback(() => {
     const handlers = {
       billingCheckoutSession: (value: string) => {
-        const billingCheckoutSession = jotaiStore.get(
+        const billingCheckoutSession = store.get(
           billingCheckoutSessionState.atom,
         );
 
@@ -27,7 +28,7 @@ export const useInitializeQueryParamState = () => {
             'requirePaymentMethod' in parsedValue &&
             !deepEqual(billingCheckoutSession, parsedValue)
           ) {
-            jotaiStore.set(
+            store.set(
               billingCheckoutSessionState.atom,
               parsedValue as BillingCheckoutSession,
             );
@@ -38,7 +39,7 @@ export const useInitializeQueryParamState = () => {
             'Failed to parse billingCheckoutSession from URL',
             error,
           );
-          jotaiStore.set(
+          store.set(
             billingCheckoutSessionState.atom,
             BILLING_CHECKOUT_SESSION_DEFAULT_VALUE,
           );
@@ -54,7 +55,7 @@ export const useInitializeQueryParamState = () => {
         handler(value);
       }
     }
-  }, []);
+  }, [store]);
 
   return { initializeQueryParamState };
 };