From f1372d16878c3fc96bae1e91dcd8155cf39b850c Mon Sep 17 00:00:00 2001
From: Alessio Gravili <alessio@gravili.de>
Date: Mon, 1 Sep 2025 20:22:42 -0700
Subject: [PATCH] fix memoization and rerendering

---
 .../RenderLexical/RichTextComponentClient.tsx |  1 -
 .../field/RenderLexical/useRenderEditor.tsx   | 71 ++++++++++---------
 .../lexical/collections/OnDemand/OnDemand.tsx | 46 +++---------
 3 files changed, 44 insertions(+), 74 deletions(-)

diff --git a/packages/richtext-lexical/src/field/RenderLexical/RichTextComponentClient.tsx b/packages/richtext-lexical/src/field/RenderLexical/RichTextComponentClient.tsx
index 7657494e42..f7a1ab8e1e 100644
--- a/packages/richtext-lexical/src/field/RenderLexical/RichTextComponentClient.tsx
+++ b/packages/richtext-lexical/src/field/RenderLexical/RichTextComponentClient.tsx
@@ -30,7 +30,6 @@ export const RichTextComponentClient: React.FC<{
   const onChange: (args: { formState: FormState; submitted?: boolean }) => Promise<FormState> =
     // eslint-disable-next-line @typescript-eslint/require-await
     React.useCallback(async ({ formState }) => {
-      console.log('updated form state', formState)
       return formState
     }, [])
 
diff --git a/packages/richtext-lexical/src/field/RenderLexical/useRenderEditor.tsx b/packages/richtext-lexical/src/field/RenderLexical/useRenderEditor.tsx
index 303205f1b8..e4ff81fcca 100644
--- a/packages/richtext-lexical/src/field/RenderLexical/useRenderEditor.tsx
+++ b/packages/richtext-lexical/src/field/RenderLexical/useRenderEditor.tsx
@@ -35,43 +35,44 @@ export const useRenderEditor_internal_ = (args: RenderLexicalServerFunctionArgs)
     void render()
   }, [serverFunction, admin, editorTarget, name, path, schemaPath, initialValue])
 
-  const WrappedComponent = React.memo(function WrappedComponent({
-    setValue,
-    value,
-  }: /**
-   * If value or setValue, or both, is provided, this component will manage its own value.
-   * If neither is passed, it will rely on the parent form to manage the value.
-   */
-  {
-    setValue?: FieldType<DefaultTypedEditorState | undefined>['setValue']
+  const WrappedComponent = React.useMemo(() => {
+    function Memoized({
+      setValue,
+      value,
+    }: /**
+     * If value or setValue, or both, is provided, this component will manage its own value.
+     * If neither is passed, it will rely on the parent form to manage the value.
+     */
+    {
+      setValue?: FieldType<DefaultTypedEditorState | undefined>['setValue']
 
-    value?: FieldType<DefaultTypedEditorState | undefined>['value']
-  }) {
-    if (!Component) {
-      return null
+      value?: FieldType<DefaultTypedEditorState | undefined>['value']
+    }) {
+      if (!Component) {
+        return null
+      }
+
+      if (typeof value === 'undefined' && !setValue) {
+        return Component
+      }
+
+      const fieldValue: FieldType<DefaultTypedEditorState | undefined> = {
+        disabled: false,
+        formInitializing: false,
+        formProcessing: false,
+        formSubmitted: false,
+        initialValue: value,
+        path: path ?? name,
+        setValue: setValue ?? (() => undefined),
+        showError: false,
+        value,
+      }
+
+      return <FieldContext value={fieldValue}>{Component}</FieldContext>
     }
-    if (typeof value === 'undefined' && !setValue) {
-      return Component
-    }
-    return (
-      <FieldContext
-        value={
-          {
-            disabled: false,
-            formInitializing: false,
-            formProcessing: false,
-            formSubmitted: false,
-            path: path ?? name,
-            setValue: setValue ?? (() => undefined),
-            showError: false,
-            value,
-          } satisfies FieldType<DefaultTypedEditorState | undefined>
-        }
-      >
-        {Component}
-      </FieldContext>
-    )
-  })
+
+    return Memoized
+  }, [Component, name, path])
 
   return { Component: WrappedComponent, renderLexical }
 }
diff --git a/test/lexical/collections/OnDemand/OnDemand.tsx b/test/lexical/collections/OnDemand/OnDemand.tsx
index 90fcc57e86..51acc721fa 100644
--- a/test/lexical/collections/OnDemand/OnDemand.tsx
+++ b/test/lexical/collections/OnDemand/OnDemand.tsx
@@ -4,16 +4,15 @@ import type { DefaultTypedEditorState } from '@payloadcms/richtext-lexical'
 import type { JSONFieldClientComponent } from 'payload'
 
 import { buildEditorState, useRenderEditor_internal_ } from '@payloadcms/richtext-lexical/client'
-import { use, useCallback, useEffect, useRef, useState } from 'react'
+import React, { useEffect, useRef, useState } from 'react'
 
 export const OnDemand: JSONFieldClientComponent = (args) => {
   const { Component, renderLexical } = useRenderEditor_internal_({
     name: 'richText',
     editorTarget: 'default',
   })
-
-  // mount the lexical runtime once
   const mounted = useRef(false)
+
   useEffect(() => {
     if (mounted.current) {
       return
@@ -22,49 +21,20 @@ export const OnDemand: JSONFieldClientComponent = (args) => {
     mounted.current = true
   }, [renderLexical])
 
-  // build the initial editor state once, with lazy init (no ref reads in render)
-  const [initialValue] = useState<DefaultTypedEditorState | undefined>(() =>
+  const [value, setValue] = useState<DefaultTypedEditorState | undefined>(() =>
     buildEditorState({ text: 'state default' }),
   )
 
-  // keep latest content in a ref so updates don’t trigger React renders
-  const latestValueRef = useRef<DefaultTypedEditorState | undefined>(initialValue)
-
-  // stable setter given to the editor; updates ref only
-  const setValueStable = useCallback((next: DefaultTypedEditorState | undefined) => {
-    // absolutely no state set here; no React re-render, no remount
-    latestValueRef.current = next
-    // if you later get access to the editor instance, this is where you'd imperatively sync it
+  const handleReset = React.useCallback(() => {
+    setValue(buildEditorState({ text: 'state default' }))
   }, [])
 
-  // If you need a "reset to default," and the editor doesn't expose an imperative API,
-  // the only reliable way is a key bump to force a remount ON RESET ONLY.
-  // This does not affect normal setValue cycles.
-  const [resetNonce, setResetNonce] = useState(0)
-  const handleReset = useCallback(() => {
-    latestValueRef.current = initialValue
-    // If you have an imperative API: editor.setEditorState(initialValue)
-    // Otherwise, remount once to guarantee visual reset:
-    setResetNonce((n) => n + 1)
-  }, [initialValue])
-
   return (
     <div>
-      <div>Default Component:</div>
-      {Component ? (
-        <Component
-          key={resetNonce}
-          // editor will call this; we won't re-render on its calls
-          setValue={setValueStable as any}
-          // initial value only; never changes so the element won’t re-render because of this prop
-          value={initialValue}
-        />
-      ) : (
-        'Loading...'
-      )}
-
+      Default Component:
+      {Component ? <Component setValue={setValue as any} value={value} /> : 'Loading...'}
       <button onClick={handleReset} style={{ marginTop: 8 }} type="button">
-        Reset to Default
+        Reset Editor State
       </button>
     </div>
   )