Refactor: Making TaskStore a regular context (#3732)

* Simplifying TaskStore by just using it as a normal context instead of using zustand. It is only really used to override things, and the pattern we used to wait until things were set in zustand clearly indicated this was the wrong approach.

* Renaming file, renaming dataElementId -> dataModelElementId

* Depth is never read(!)

* Suggestions from coderabbit in review

* I found another issue. The useLayoutSetId() hook was overly complicated, and did much of the same stuff that useCurrentLayoutSetId() already did. The main difference was the taskId source, where useLayoutSetId() preferred the one from the URL, so I renamed it to reflect that.

* Fixing import

---------

Co-authored-by: Ole Martin Handeland <[email protected]>

Author: olemartinorg

src/core/contexts/TaskOverrides.tsx

@@ -0,0 +1,32 @@
+import React, { createContext, useContext } from 'react';
+import type { PropsWithChildren } from 'react';
+
+interface TaskOverridesContext {
+  taskId?: string;
+  dataModelType?: string;
+  dataModelElementId?: string;
+  layoutSetId?: string;
+}
+
+const Context = createContext<TaskOverridesContext>({});
+Context.displayName = 'TaskOverridesContext';
+
+type Props = PropsWithChildren & TaskOverridesContext;
+export function TaskOverrides({ children, ...overrides }: Props) {
+  const parentContext = useContext(Context);
+
+  return (
+    <Context.Provider
+      value={{
+        taskId: overrides.taskId ?? parentContext.taskId,
+        dataModelType: overrides.dataModelType ?? parentContext.dataModelType,
+        dataModelElementId: overrides.dataModelElementId ?? parentContext.dataModelElementId,
+        layoutSetId: overrides.layoutSetId ?? parentContext.layoutSetId,
+      }}
+    >
+      {children}
+    </Context.Provider>
+  );
+}
+
+export const useTaskOverrides = () => useContext(Context);

src/core/contexts/taskStoreContext.tsx

@@ -2,7 +2,7 @@ import React, { useEffect } from 'react';
 
 import deepEqual from 'fast-deep-equal';
 
-import { useTaskStore } from 'src/core/contexts/taskStoreContext';
+import { useTaskOverrides } from 'src/core/contexts/TaskOverrides';
 import { useApplicationMetadata } from 'src/features/applicationMetadata/ApplicationMetadataProvider';
 import { isAttachmentUploaded } from 'src/features/attachments/index';
 import { DEFAULT_DEBOUNCE_TIMEOUT } from 'src/features/formData/types';
@@ -89,15 +89,14 @@ function useNodeAttachments(): AttachmentRecord {
   const { indexedId, baseId } = parent;
   const nodeData = useFormDataFor(baseId) as IComponentFormData<CompWithBehavior<'canHaveAttachments'>>;
 
-  const overriddenTaskId = useTaskStore((state) => state.overriddenTaskId);
+  const overriddenTaskId = useTaskOverrides()?.taskId;
 
   const application = useApplicationMetadata();
   const currentTask = useProcessQuery().data?.currentTask?.elementId;
   const data = useInstanceDataElements(baseId);
 
   const mappedAttachments = useMemoDeepEqual(() => {
-    const taskId = overriddenTaskId ? overriddenTaskId : currentTask;
-
+    const taskId = overriddenTaskId ?? currentTask;
     return mapAttachments(indexedId, baseId, data, application, taskId, nodeData);
   }, [indexedId, baseId, data, application, currentTask, nodeData, overriddenTaskId]);
 

src/features/attachments/StoreAttachmentsInNode.tsx

@@ -6,7 +6,7 @@ import deepEqual from 'fast-deep-equal';
 import { createStore } from 'zustand';
 import type { JSONSchema7 } from 'json-schema';
 
-import { useTaskStore } from 'src/core/contexts/taskStoreContext';
+import { useTaskOverrides } from 'src/core/contexts/TaskOverrides';
 import { createZustandContext } from 'src/core/contexts/zustandContext';
 import { DisplayError } from 'src/core/errorHandling/DisplayError';
 import { Loader } from 'src/core/loading/Loader';
@@ -170,8 +170,9 @@ function DataModelsLoader() {
   const layoutSetId = useCurrentLayoutSetId();
 
   // Subform
-  const overriddenDataElementId = useTaskStore((state) => state.overriddenDataElementId);
-  const overriddenDataType = useTaskStore((state) => state.overriddenDataModelType);
+  const overrides = useTaskOverrides();
+  const overriddenDataElementId = overrides?.dataModelElementId;
+  const overriddenDataType = overrides?.dataModelType;
 
   // Find all data types referenced in dataModelBindings in the layout
   useEffect(() => {

src/features/datamodel/DataModelsProvider.tsx

@@ -2,7 +2,7 @@ import { useCallback, useMemo } from 'react';
 
 import type { JSONSchema7 } from 'json-schema';
 
-import { useTaskStore } from 'src/core/contexts/taskStoreContext';
+import { useTaskOverrides } from 'src/core/contexts/TaskOverrides';
 import { useApplicationMetadata } from 'src/features/applicationMetadata/ApplicationMetadataProvider';
 import {
   getCurrentDataTypeForApplication,
@@ -34,7 +34,7 @@ export function useCurrentDataModelDataElementId() {
   const layoutSets = useLayoutSets();
   const taskId = useProcessTaskId();
 
-  const overriddenDataElementId = useTaskStore((s) => s.overriddenDataElementId);
+  const overriddenDataElementId = useTaskOverrides()?.dataModelElementId;
 
   // Instance data elements will update often (after each save), so we have to use a selector to make
   // sure components don't re-render too often.
@@ -130,7 +130,7 @@ export function useDataModelUrl({ dataType, dataElementId, language, prefillFrom
 }
 
 export function useCurrentDataModelName() {
-  const overriddenDataModelType = useTaskStore((state) => state.overriddenDataModelType);
+  const overriddenDataModelType = useTaskOverrides()?.dataModelType;
 
   const application = useApplicationMetadata();
   const layoutSets = useLayoutSets();

src/features/datamodel/useBindingSchema.tsx

@@ -10,7 +10,8 @@ import { LayoutInspectorItem } from 'src/features/devtools/components/LayoutInsp
 import { SplitView } from 'src/features/devtools/components/SplitView/SplitView';
 import { useDevToolsStore } from 'src/features/devtools/data/DevToolsStore';
 import { useLayoutValidationForPage } from 'src/features/devtools/layoutValidation/useLayoutValidation';
-import { useLayouts, useLayoutSetId } from 'src/features/form/layout/LayoutsContext';
+import { useLayouts } from 'src/features/form/layout/LayoutsContext';
+import { useLayoutSetIdFromUrl } from 'src/features/form/layoutSets/useCurrentLayoutSet';
 import { useCurrentView } from 'src/hooks/useNavigatePage';
 import { parseAndCleanText } from 'src/language/sharedLanguage';
 import type { LayoutContextValue } from 'src/features/form/layout/LayoutsContext';
@@ -20,7 +21,7 @@ export const LayoutInspector = () => {
   const setSelectedComponent = useDevToolsStore((state) => state.actions.layoutInspectorSet);
   const currentView = useCurrentView();
   const layouts = useLayouts();
-  const currentLayoutSetId = useLayoutSetId();
+  const currentLayoutSetId = useLayoutSetIdFromUrl();
   const [componentProperties, setComponentProperties] = useState<string | null>(null);
   const [propertiesHaveChanged, setPropertiesHaveChanged] = useState(false);
   const [error, setError] = useState<boolean>(false);

src/features/devtools/components/LayoutInspector/LayoutInspector.tsx

@@ -6,16 +6,14 @@ import { useAppQueries } from 'src/core/contexts/AppQueriesProvider';
 import { ContextNotProvided } from 'src/core/contexts/context';
 import { delayedContext } from 'src/core/contexts/delayedContext';
 import { createQueryContext } from 'src/core/contexts/queryContext';
-import { useTaskStore } from 'src/core/contexts/taskStoreContext';
 import { useCurrentDataModelName } from 'src/features/datamodel/useBindingSchema';
 import { cleanLayout } from 'src/features/form/layout/cleanLayout';
 import { makeLayoutLookups } from 'src/features/form/layout/makeLayoutLookups';
 import { applyLayoutQuirks } from 'src/features/form/layout/quirks';
 import { useLayoutSets } from 'src/features/form/layoutSets/LayoutSetsProvider';
-import { useCurrentLayoutSetId } from 'src/features/form/layoutSets/useCurrentLayoutSet';
+import { useLayoutSetIdFromUrl } from 'src/features/form/layoutSets/useCurrentLayoutSet';
 import { useInstanceDataQuery } from 'src/features/instance/InstanceContext';
 import { useProcessQuery } from 'src/features/instance/useProcessQuery';
-import { useNavigationParam } from 'src/hooks/navigation';
 import { makeLikertChildId } from 'src/layout/Likert/Generator/makeLikertChildId';
 import type { QueryDefinition } from 'src/core/queries/usePrefetchQuery';
 import type { CompExternal, ILayoutCollection, ILayouts } from 'src/layout/layout';
@@ -45,7 +43,7 @@ export function useLayoutQueryDef(
 
 function useLayoutQuery() {
   const { data: process } = useProcessQuery();
-  const currentLayoutSetId = useLayoutSetId();
+  const currentLayoutSetId = useLayoutSetIdFromUrl();
   const defaultDataModel = useCurrentDataModelName() ?? 'unknown';
   const hasInstance = !!useInstanceDataQuery().data;
 
@@ -78,30 +76,6 @@ const { Provider, useCtx, useLaxCtx } = delayedContext(() =>
   }),
 );
 
-export function useLayoutSetId() {
-  const layoutSets = useLayoutSets();
-  const currentProcessLayoutSetId = useCurrentLayoutSetId();
-  const taskId = useNavigationParam('taskId');
-
-  const overriddenLayoutSetId = useTaskStore((state) => state.overriddenLayoutSetId);
-
-  if (overriddenLayoutSetId) {
-    return overriddenLayoutSetId;
-  }
-
-  const layoutSetId =
-    taskId != null
-      ? layoutSets.find((set) => {
-          if (set.tasks?.length) {
-            return set.tasks.includes(taskId);
-          }
-          return false;
-        })?.id
-      : undefined;
-
-  return layoutSetId ?? currentProcessLayoutSetId;
-}
-
 export function useDataTypeFromLayoutSet(layoutSetName: string | undefined) {
   const layoutSets = useLayoutSets();
   return layoutSets.find((set) => set.id === layoutSetName)?.dataType;

src/features/form/layout/LayoutsContext.tsx

@@ -1,19 +1,31 @@
 import { ContextNotProvided } from 'src/core/contexts/context';
-import { useTaskStore } from 'src/core/contexts/taskStoreContext';
+import { useTaskOverrides } from 'src/core/contexts/TaskOverrides';
 import { useLaxApplicationMetadata } from 'src/features/applicationMetadata/ApplicationMetadataProvider';
 import { getCurrentLayoutSet } from 'src/features/applicationMetadata/appMetadataUtils';
 import { useLaxLayoutSets } from 'src/features/form/layoutSets/LayoutSetsProvider';
 import { useProcessTaskId } from 'src/features/instance/useProcessTaskId';
+import { useNavigationParam } from 'src/hooks/navigation';
 
-export function useCurrentLayoutSetId() {
-  return useCurrentLayoutSet()?.id;
+/**
+ * This is a variant that prefers the taskId from the URL. The alternative useCurrentLayoutSetId() and
+ * useCurrentLayoutSet() will prefer the taskId from the current process state (i.e., where the process is right now,
+ * not necessarily what the user is looking at right now).
+ */
+export function useLayoutSetIdFromUrl() {
+  const taskId = useNavigationParam('taskId');
+  return useCurrentLayoutSetId(taskId);
 }
 
-export function useCurrentLayoutSet() {
+export function useCurrentLayoutSetId(taskId?: string) {
+  return useCurrentLayoutSet(taskId)?.id;
+}
+
+export function useCurrentLayoutSet(_taskId?: string) {
   const application = useLaxApplicationMetadata();
   const layoutSets = useLaxLayoutSets();
-  const taskId = useProcessTaskId();
-  const overriddenLayoutSetId = useTaskStore((state) => state.overriddenLayoutSetId);
+  const processTaskId = useProcessTaskId();
+  const taskId = _taskId ?? processTaskId;
+  const overriddenLayoutSetId = useTaskOverrides()?.layoutSetId;
 
   if (application === ContextNotProvided || layoutSets === ContextNotProvided) {
     return undefined;

src/features/form/layoutSets/useCurrentLayoutSet.ts

@@ -7,8 +7,8 @@ import { useAppQueries } from 'src/core/contexts/AppQueriesProvider';
 import { ContextNotProvided } from 'src/core/contexts/context';
 import { delayedContext } from 'src/core/contexts/delayedContext';
 import { createQueryContext } from 'src/core/contexts/queryContext';
-import { useLayoutSetId } from 'src/features/form/layout/LayoutsContext';
 import { useLaxGlobalUISettings } from 'src/features/form/layoutSets/LayoutSetsProvider';
+import { useLayoutSetIdFromUrl } from 'src/features/form/layoutSets/useCurrentLayoutSet';
 import { useShallowMemo } from 'src/hooks/useShallowMemo';
 import type { QueryDefinition } from 'src/core/queries/usePrefetchQuery';
 import type { GlobalPageSettings, ILayoutSettings, NavigationPageGroup } from 'src/layout/common.generated';
@@ -23,7 +23,7 @@ export function useLayoutSettingsQueryDef(layoutSetId?: string): QueryDefinition
 }
 
 function useLayoutSettingsQuery() {
-  const layoutSetId = useLayoutSetId();
+  const layoutSetId = useLayoutSetIdFromUrl();
   const query = useQuery(useLayoutSettingsQueryDef(layoutSetId));
 
   useEffect(() => {

src/features/form/layoutSettings/LayoutSettingsContext.tsx

@@ -1,9 +1,9 @@
-import { useTaskStore } from 'src/core/contexts/taskStoreContext';
+import { useTaskOverrides } from 'src/core/contexts/TaskOverrides';
 import { useProcessQuery } from 'src/features/instance/useProcessQuery';
 import { useNavigationParam } from 'src/hooks/navigation';
 
 export function useProcessTaskId() {
-  const overriddenTaskId = useTaskStore((state) => state.overriddenTaskId);
+  const overriddenTaskId = useTaskOverrides()?.taskId;
   const processTaskId = useProcessQuery().data?.currentTask?.elementId;
   const urlTaskId = useNavigationParam('taskId');
   return overriddenTaskId ?? processTaskId ?? urlTaskId;