Zustand re-rendering issue after restoring the state from indexedDB

[anni1236012]

I am new to zustand and I am running into a re-rendering issue and it only happens when state is restored from the indexedDB. Which means only on page refresh. I've spent a week figuring out but no luck. If I don't refresh the page then it correctly re-renders on the update. I feel something to do with produce method imported from immer but I am not able to pinpoint the issue.

Here is an example. In updateNodeText function, I am updating the text but it does not re-render.

 if (state.selectedNode.node.type === "JSXText") {
         state.selectedNode.node.value = text;
  }

For simplicity, I've removed the logic from the rest of the functions.

import { create } from "zustand";
import { persist } from "zustand/middleware";
import { produce } from "immer";
import {
  EditorState,
  Page,
  Component,
  SelectedNode,
  PrebuiltComponent,
  Project,
} from "@/types";
import { useHydration } from "./use-hydration";
import { createJSONStorage } from "zustand/middleware";

// Prebuilt components data
const prebuiltComponents: PrebuiltComponent[] = [
  {
    name: "Deep",
    path: "/prebuilt-components/deep.tsx",
    description: "Deep component",
  },
];

const defaultPage: Page = {
  id: crypto.randomUUID(),
  name: "Home",
  components: [],
  isDefault: true,
};

// Add IndexedDB storage implementation
const createIndexedDBStorage = () => {
  const dbName = "editor-db";
  const storeName = "editor-store";
  const version = 1;

  let dbInstance: IDBDatabase | null = null;

  const getDB = async (): Promise<IDBDatabase> => {
    // Return existing connection if available
    if (dbInstance) return dbInstance;

    return new Promise((resolve, reject) => {
      const request = indexedDB.open(dbName, version);

      request.onerror = () => {
        console.error("Failed to open IndexedDB:", request.error);
        reject(request.error);
      };

      request.onsuccess = () => {
        dbInstance = request.result;

        // Handle connection closing
        dbInstance.onclose = () => {
          dbInstance = null;
        };

        // Handle connection errors
        dbInstance.onerror = (event) => {
          console.error("IndexedDB error:", event);
        };

        resolve(dbInstance);
      };

      request.onupgradeneeded = (event) => {
        const db = (event.target as IDBOpenDBRequest).result;
        if (!db.objectStoreNames.contains(storeName)) {
          db.createObjectStore(storeName);
        }
      };
    });
  };

  const storage = {
    getItem: async (name: string): Promise<any> => {
      const db = await getDB();
      return new Promise((resolve, reject) => {
        try {
          const transaction = db.transaction(storeName, "readonly");
          const store = transaction.objectStore(storeName);
          const request = store.get(name);

          request.onerror = () => {
            console.error("Error reading from IndexedDB:", request.error);
            reject(request.error);
          };

          request.onsuccess = () => {
            resolve(request.result || null);
          };

          // Handle transaction completion
          transaction.oncomplete = () => {
            // transaction complete
          };
        } catch (error) {
          console.error("Transaction error:", error);
          reject(error);
        }
      });
    },

    setItem: async (name: string, value: any): Promise<void> => {
      const db = await getDB();
      return new Promise((resolve, reject) => {
        try {
          const transaction = db.transaction(storeName, "readwrite");
          const store = transaction.objectStore(storeName);
          const request = store.put(value, name);

          request.onerror = () => {
            console.error("Error writing to IndexedDB:", request.error);
            reject(request.error);
          };

          request.onsuccess = () => {
            resolve();
          };

          // Handle transaction completion
          transaction.oncomplete = () => {
            // transaction complete
          };
        } catch (error) {
          console.error("Transaction error:", error);
          reject(error);
        }
      });
    },

    removeItem: async (name: string): Promise<void> => {
      const db = await getDB();
      return new Promise((resolve, reject) => {
        try {
          const transaction = db.transaction(storeName, "readwrite");
          const store = transaction.objectStore(storeName);
          const request = store.delete(name);

          request.onerror = () => {
            console.error("Error removing from IndexedDB:", request.error);
            reject(request.error);
          };

          request.onsuccess = () => {
            resolve();
          };

          // Handle transaction completion
          transaction.oncomplete = () => {
            // transaction complete
          };
        } catch (error) {
          console.error("Transaction error:", error);
          reject(error);
        }
      });
    },
  };

  return storage;
};

export const useEditor = create(
  persist(
    (set, get) => ({
      // Initial state
      selectedComponentId: null,
      currentProjectId: null,
      selectedNode: null,
      prebuiltComponents,

      projects: [],
      // New actions
      createProject: (name: string, description?: string) =>
        set(
          produce((state) => {
            const newProject: Project = {
              id: crypto.randomUUID(),
              name,
              description,
              pages: [defaultPage],
              createdAt: Date.now(),
              updatedAt: Date.now(),
              previewImage: "",
              preview: false,
              selectedPageId: defaultPage.id,
            };
            state.projects.push(newProject);
            state.currentProjectId = newProject.id;
            state.selectedComponentId = null;
            state.selectedNode = null;
          })
        ),

      // Node actions
      setSelectedNode: (node: SelectedNode) =>
        set(
          produce((state) => {
            const currentProject = state.projects.find(
              (p: Project) => p.id === state.currentProjectId
            );
            if (!currentProject) return;

            state.selectedNode = node;
            state.selectedPageId = node.pageId;
            state.selectedComponentId = node.componentId;
          })
        ),

      updateNodeText: (text: string) =>
        set(
          produce((state) => {
            if (!state.selectedNode?.node) return;

            const currentProject = state.projects.find(
              (p: Project) => p.id === state.currentProjectId
            );
            if (!currentProject) return;

            if (state.selectedNode.node.type === "JSXText") {
              state.selectedNode.node.value = text;
            }
          })
        ),
    }),
    {
      name: "editor-storage",
      storage: createJSONStorage(() => createIndexedDBStorage()),
      partialize: (state: EditorState) => ({
        projects: state.projects,
        selectedComponentId: state.selectedComponentId,
        currentProjectId: state.currentProjectId,
      }),
      onRehydrateStorage: () => (state) => {
        if (!state) return;

        useHydration.setState({ hydrated: true });
      },
    }
  )
);

// Create a wrapper hook that ensures hydration
export function useHydratedEditor<T>(
  selector: (state: EditorState) => T,
  defaultValue: T
): T {
  const hydrated = useHydration((state) => state.hydrated);
  const value = useEditor(selector);

  return hydrated ? value : defaultValue;
}

[anni1236012]

I figured out the issue. The problem occurs because when the state is rehydrated from IndexedDB, new object references are created for the AST nodes, breaking the reference equality between selectedNode.node and the nodes in the component tree. instead of using the huge object as it is, I had to recreate the object by using the raw component string.

In short, issue was due to object references were not the same after restored the state from the database.

[anni1236012]

Here is the video demo of the problem I was facing if it is not clear what I explained above.
https://x.com/i/status/1880473541482852438