perf(richtext-lexical)!: significantly reduce lexical rerendering and amount of network requests from blocks (#9255)

The field RSC now provides an initial state for all lexical blocks. This completely obliterates any flashes and lexical block loading states when loading or saving a document. Previously, when a document is loaded or saved, every lexical block was sending a network request in order to fetch their form state. Now, this is batched and handled in the lexical server component. All lexical block form states are sent to the client together with the parent lexical field, and are thus available immediately. We also do the same with block collapsed preferences. Thus, there are no loading states or layout shifts/flashes of blocks anymore. Additionally, when saving a document while your cursor is inside a lexical field, the cursor position is preserved. Previously, a document save would kick your cursor out of the lexical field. ## Look at how nice this is: https://github.com/user-attachments/assets/21d736d4-8f80-4df0-a782-7509edd993da **BREAKING:** This removes the `feature.hooks.load` and `feature.hooks.save` interfaces from custom lexical features, as they weren't used internally and added unnecessary, additional overhead. If you have custom features that use those, you can migrate to using normal payload hooks that run on the server instead of the client.
2024-11-17 01:31:55 -07:00
parent abe4cc87ca
commit 35917c67d7
20 changed files with 394 additions and 302 deletions
--- a/packages/payload/src/admin/forms/Field.ts
+++ b/packages/payload/src/admin/forms/Field.ts
@@ -3,7 +3,8 @@ import type { MarkOptional } from 'ts-essentials'

 import type { SanitizedFieldPermissions, User } from '../../auth/types.js'
 import type { ClientBlock, ClientField, Field } from '../../fields/config/types.js'
-import type { Payload } from '../../types/index.js'
+import type { DocumentPreferences } from '../../preferences/types.js'
+import type { Operation, Payload, PayloadRequest } from '../../types/index.js'
 import type {
  ClientTab,
  Data,
@@ -66,6 +67,7 @@ export type FieldPaths = {

 export type ServerComponentProps = {
  clientField: ClientFieldWithOptionalType
+  collectionSlug: string
  data: Data
  field: Field
  /**
@@ -78,8 +80,12 @@ export type ServerComponentProps = {
   */
  formState: FormState
  i18n: I18nClient
+  id?: number | string
+  operation: Operation
  payload: Payload
  permissions: SanitizedFieldPermissions
+  preferences: DocumentPreferences
+  req: PayloadRequest
  siblingData: Data
  user: User
 }
--- a/packages/richtext-lexical/src/features/blocks/client/component/BlockContent.tsx
+++ b/packages/richtext-lexical/src/features/blocks/client/component/BlockContent.tsx
@@ -1,13 +1,9 @@
 'use client'
 import type { ClientField, FormState } from 'payload'

-import { RenderFields } from '@payloadcms/ui'
+import { RenderFields, useFormSubmitted } from '@payloadcms/ui'
 import React, { createContext, useMemo } from 'react'

-import type { BlockFields } from '../../server/nodes/BlocksNode.js'
-
-import { useFormSave } from './FormSavePlugin.js'
-
 type Props = {
  baseClass: string
  BlockDrawer: React.FC
@@ -24,11 +20,11 @@ type Props = {
  }>
  CustomBlock: React.ReactNode
  EditButton: React.FC
-  formData: BlockFields
+  errorCount: number
  formSchema: ClientField[]
  initialState: false | FormState | undefined
-  nodeKey: string

+  nodeKey: string
  RemoveButton: React.FC
 }

@@ -66,14 +62,16 @@ export const BlockContent: React.FC<Props> = (props) => {
    Collapsible,
    CustomBlock,
    EditButton,
-    formData,
+    errorCount,
    formSchema,
    initialState,
    nodeKey,
    RemoveButton,
  } = props

-  const { errorCount, fieldHasErrors } = useFormSave({ disabled: !initialState, formData, nodeKey })
+  const hasSubmitted = useFormSubmitted()
+
+  const fieldHasErrors = hasSubmitted && errorCount > 0

  const CollapsibleWithErrorProps = useMemo(
    () =>
--- a/packages/richtext-lexical/src/features/blocks/client/component/FormSavePlugin.tsx
+++ b/packages/richtext-lexical/src/features/blocks/client/component/FormSavePlugin.tsx
@@ -1,115 +0,0 @@
-'use client'
-import type { FormState } from 'payload'
-
-import { useLexicalComposerContext } from '@lexical/react/LexicalComposerContext'
-import { useAllFormFields, useFormSubmitted } from '@payloadcms/ui'
-import { dequal } from 'dequal/lite'
-import { $getNodeByKey } from 'lexical'
-import { reduceFieldsToValues } from 'payload/shared'
-import React, { useCallback, useEffect } from 'react'
-
-import type { BlockFields } from '../../server/nodes/BlocksNode.js'
-
-import { $isBlockNode } from '../nodes/BlocksNode.js'
-import { removeEmptyArrayValues } from './removeEmptyArrayValues.js'
-
-type Props = {
-  disabled?: boolean
-  formData: BlockFields
-  nodeKey: string
-}
-
-// Recursively remove all undefined values from even being present in formData, as they will
-// cause isDeepEqual to return false if, for example, formData has a key that fields.data
-// does not have, even if it's undefined.
-// Currently, this happens if a block has another sub-blocks field. Inside formData, that sub-blocks field has an undefined blockName property.
-// Inside of fields.data however, that sub-blocks blockName property does not exist at all.
-function removeUndefinedAndNullAndEmptyArraysRecursively(obj: object) {
-  for (const key in obj) {
-    const value = obj[key]
-    if (Array.isArray(value) && !value?.length) {
-      delete obj[key]
-    } else if (value && typeof value === 'object') {
-      removeUndefinedAndNullAndEmptyArraysRecursively(value)
-    } else if (value === undefined || value === null) {
-      delete obj[key]
-    }
-  }
-}
-
-export const useFormSave = (props: Props): { errorCount: number; fieldHasErrors: boolean } => {
-  const { disabled, formData, nodeKey } = props
-
-  const [_fields] = useAllFormFields()
-
-  const fields = removeEmptyArrayValues({ fields: _fields })
-
-  // Pass in fields, and indicate if you'd like to "unflatten" field data.
-  // The result below will reflect the data stored in the form at the given time
-  const newFormData = reduceFieldsToValues(fields, true)
-
-  const [editor] = useLexicalComposerContext()
-
-  const hasSubmitted = useFormSubmitted()
-
-  const [errorCount, setErrorCount] = React.useState(0)
-
-  const fieldHasErrors = hasSubmitted && errorCount > 0
-
-  const onFormChange = useCallback(
-    ({
-      fullFieldsWithValues,
-      newFormData,
-    }: {
-      fullFieldsWithValues: FormState
-      newFormData: BlockFields
-    }) => {
-      newFormData.id = formData.id
-      newFormData.blockType = formData.blockType
-
-      removeUndefinedAndNullAndEmptyArraysRecursively(newFormData)
-      removeUndefinedAndNullAndEmptyArraysRecursively(formData)
-
-      // Only update if the data has actually changed. Otherwise, we may be triggering an unnecessary value change,
-      // which would trigger the "Leave without saving" dialog unnecessarily
-      if (!dequal(formData, newFormData)) {
-        // Running this in the next tick in the meantime fixes this issue: https://github.com/payloadcms/payload/issues/4108
-        // I don't know why. When this is called immediately, it might focus out of a nested lexical editor field if an update is made there.
-        // My hypothesis is that the nested editor might not have fully finished its update cycle yet. By updating in the next tick, we
-        // ensure that the nested editor has finished its update cycle before we update the block node.
-        setTimeout(() => {
-          editor.update(() => {
-            const node = $getNodeByKey(nodeKey)
-            if (node && $isBlockNode(node)) {
-              node.setFields(newFormData)
-            }
-          })
-        }, 0)
-      }
-
-      // update error count
-      if (hasSubmitted) {
-        let rowErrorCount = 0
-        for (const formField of Object.values(fullFieldsWithValues)) {
-          if (formField?.valid === false) {
-            rowErrorCount++
-          }
-        }
-        setErrorCount(rowErrorCount)
-      }
-    },
-    [editor, nodeKey, hasSubmitted, formData],
-  )
-
-  useEffect(() => {
-    if (disabled) {
-      return
-    }
-    onFormChange({ fullFieldsWithValues: fields, newFormData: newFormData as BlockFields })
-  }, [newFormData, fields, onFormChange, disabled])
-
-  return {
-    errorCount,
-    fieldHasErrors,
-  }
-}
--- a/packages/richtext-lexical/src/features/blocks/client/component/index.tsx
+++ b/packages/richtext-lexical/src/features/blocks/client/component/index.tsx
@@ -19,8 +19,8 @@ import {
  useTranslation,
 } from '@payloadcms/ui'
 import { abortAndIgnore } from '@payloadcms/ui/shared'
-import { reduceFieldsToValues } from 'payload/shared'
-import React, { useCallback, useEffect, useMemo, useRef, useState } from 'react'
+import { deepCopyObjectSimpleWithoutReactComponents, reduceFieldsToValues } from 'payload/shared'
+import React, { useCallback, useEffect, useMemo, useRef } from 'react'

 const baseClass = 'lexical-block'
 import { useLexicalComposerContext } from '@lexical/react/LexicalComposerContext'
@@ -36,6 +36,7 @@ import { useLexicalDrawer } from '../../../../utilities/fieldsDrawer/useLexicalD
 import { $isBlockNode } from '../nodes/BlocksNode.js'
 import { BlockContent } from './BlockContent.js'
 import './index.scss'
+import { removeEmptyArrayValues } from './removeEmptyArrayValues.js'

 type Props = {
  readonly children?: React.ReactNode
@@ -51,6 +52,7 @@ export const BlockComponent: React.FC<Props> = (props) => {
    fieldProps: {
      featureClientSchemaMap,
      field: parentLexicalRichTextField,
+      initialLexicalFormState,
      permissions,
      readOnly,
      schemaPath,
@@ -59,6 +61,7 @@ export const BlockComponent: React.FC<Props> = (props) => {
  } = useEditorConfigContext()
  const onChangeAbortControllerRef = useRef(new AbortController())
  const editDepth = useEditDepth()
+  const [errorCount, setErrorCount] = React.useState(0)

  const drawerSlug = formatDrawerSlug({
    slug: `lexical-blocks-create-${uuidFromContext}-${formData.id}`,
@@ -69,32 +72,36 @@ export const BlockComponent: React.FC<Props> = (props) => {
  // Used for saving collapsed to preferences (and gettin' it from there again)
  // Remember, these preferences are scoped to the whole document, not just this form. This
  // is important to consider for the data path used in setDocFieldPreferences
-  const { docPermissions, getDocPreferences, setDocFieldPreferences } = useDocumentInfo()
+  const { getDocPreferences, setDocFieldPreferences } = useDocumentInfo()
  const [editor] = useLexicalComposerContext()

  const { getFormState } = useServerFunctions()
-
-  const [initialState, setInitialState] = useState<false | FormState | undefined>(false)
-
  const schemaFieldsPath = `${schemaPath}.lexical_internal_feature.blocks.lexical_blocks.${formData.blockType}.fields`

-  const componentMapRenderedBlockPath = `${schemaPath}.lexical_internal_feature.blocks.lexical_blocks.${formData.blockType}`
+  const [initialState, setInitialState] = React.useState<false | FormState | undefined>(
+    initialLexicalFormState?.[formData.id]?.formState
+      ? {
+          ...initialLexicalFormState?.[formData.id]?.formState,
+          blockName: {
+            initialValue: formData.blockName,
+            passesCondition: true,
+            valid: true,
+            value: formData.blockName,
+          },
+        }
+      : false,
+  )

-  const clientSchemaMap = featureClientSchemaMap['blocks']
-
-  const blocksField: BlocksFieldClient = clientSchemaMap[
-    componentMapRenderedBlockPath
-  ][0] as BlocksFieldClient
-
-  const clientBlock = blocksField.blocks[0]
-
-  const { i18n, t } = useTranslation<object, string>()
-
-  // Field Schema
+  // Initial state for newly created blocks
  useEffect(() => {
    const abortController = new AbortController()

    const awaitInitialState = async () => {
+      /*
+       * This will only run if a new block is created. For all existing blocks that are loaded when the document is loaded, or when the form is saved,
+       * this is not run, as the lexical field RSC will fetch the state server-side and pass it to the client. That way, we avoid unnecessary client-side
+       * requests. Though for newly created blocks, we need to fetch the state client-side, as the server doesn't know about the block yet.
+       */
      const { state } = await getFormState({
        id,
        collectionSlug,
@@ -110,7 +117,7 @@ export const BlockComponent: React.FC<Props> = (props) => {

      if (state) {
        state.blockName = {
-          initialValue: '',
+          initialValue: formData.blockName,
          passesCondition: true,
          valid: true,
          value: formData.blockName,
@@ -120,7 +127,7 @@ export const BlockComponent: React.FC<Props> = (props) => {
      }
    }

-    if (formData) {
+    if (formData && !initialState) {
      void awaitInitialState()
    }

@@ -131,14 +138,31 @@ export const BlockComponent: React.FC<Props> = (props) => {
    getFormState,
    schemaFieldsPath,
    id,
+    formData,
+    initialState,
    collectionSlug,
    globalSlug,
    getDocPreferences,
-    // DO NOT ADD FORMDATA HERE! Adding formData will kick you out of sub block editors while writing.
  ])

+  const [isCollapsed, setIsCollapsed] = React.useState<boolean>(
+    initialLexicalFormState?.[formData.id]?.collapsed ?? false,
+  )
+
+  const componentMapRenderedBlockPath = `${schemaPath}.lexical_internal_feature.blocks.lexical_blocks.${formData.blockType}`
+
+  const clientSchemaMap = featureClientSchemaMap['blocks']
+
+  const blocksField: BlocksFieldClient = clientSchemaMap[
+    componentMapRenderedBlockPath
+  ][0] as BlocksFieldClient
+
+  const clientBlock = blocksField.blocks[0]
+
+  const { i18n, t } = useTranslation<object, string>()
+
  const onChange = useCallback(
-    async ({ formState: prevFormState }) => {
+    async ({ formState: prevFormState, submit }) => {
      abortAndIgnore(onChangeAbortControllerRef.current)

      const controller = new AbortController()
@@ -147,7 +171,9 @@ export const BlockComponent: React.FC<Props> = (props) => {
      const { state: newFormState } = await getFormState({
        id,
        collectionSlug,
-        docPermissions,
+        docPermissions: {
+          fields: true,
+        },
        docPreferences: await getDocPreferences(),
        formState: prevFormState,
        globalSlug,
@@ -160,11 +186,36 @@ export const BlockComponent: React.FC<Props> = (props) => {
        return prevFormState
      }

-      newFormState.blockName = {
-        initialValue: '',
-        passesCondition: true,
-        valid: true,
-        value: formData.blockName,
+      newFormState.blockName = prevFormState.blockName
+
+      const newFormStateData: BlockFields = reduceFieldsToValues(
+        removeEmptyArrayValues({
+          fields: deepCopyObjectSimpleWithoutReactComponents(newFormState),
+        }),
+        true,
+      ) as BlockFields
+
+      setTimeout(() => {
+        editor.update(() => {
+          const node = $getNodeByKey(nodeKey)
+          if (node && $isBlockNode(node)) {
+            const newData = {
+              ...newFormStateData,
+              blockType: formData.blockType,
+            }
+            node.setFields(newData)
+          }
+        })
+      }, 0)
+
+      if (submit) {
+        let rowErrorCount = 0
+        for (const formField of Object.values(newFormState)) {
+          if (formField?.valid === false) {
+            rowErrorCount++
+          }
+        }
+        setErrorCount(rowErrorCount)
      }

      return newFormState
@@ -174,11 +225,12 @@ export const BlockComponent: React.FC<Props> = (props) => {
      getFormState,
      id,
      collectionSlug,
-      docPermissions,
      getDocPreferences,
      globalSlug,
      schemaFieldsPath,
-      formData.blockName,
+      formData.blockType,
+      editor,
+      nodeKey,
    ],
  )

@@ -201,17 +253,6 @@ export const BlockComponent: React.FC<Props> = (props) => {
    ? getTranslation(clientBlock.labels.singular, i18n)
    : clientBlock?.slug

-  const [isCollapsed, setIsCollapsed] = React.useState<boolean>()
-
-  // TODO: Load collapsed preferences in RSC Field component
-  useEffect(() => {
-    void getDocPreferences().then((currentDocPreferences) => {
-      const currentFieldPreferences = currentDocPreferences?.fields[parentLexicalRichTextField.name]
-      const collapsedArray = currentFieldPreferences?.collapsed
-      setIsCollapsed(collapsedArray ? collapsedArray.includes(formData.id) : false)
-    })
-  }, [parentLexicalRichTextField.name, formData.id, getDocPreferences])
-
  const onCollapsedChange = useCallback(
    (changedCollapsed: boolean) => {
      void getDocPreferences().then((currentDocPreferences) => {
@@ -408,7 +449,11 @@ export const BlockComponent: React.FC<Props> = (props) => {
    }
    return (
      <Form
-        beforeSubmit={[onChange]}
+        beforeSubmit={[
+          async ({ formState }) => {
+            return await onChange({ formState, submit: true })
+          },
+        ]}
        fields={clientBlock.fields}
        initialState={initialState}
        onChange={[onChange]}
@@ -433,7 +478,7 @@ export const BlockComponent: React.FC<Props> = (props) => {
          Collapsible={BlockCollapsible}
          CustomBlock={CustomBlock}
          EditButton={EditButton}
-          formData={formData}
+          errorCount={errorCount}
          formSchema={clientBlock.fields}
          initialState={initialState}
          nodeKey={nodeKey}
@@ -448,6 +493,7 @@ export const BlockComponent: React.FC<Props> = (props) => {
    RemoveButton,
    EditButton,
    editor,
+    errorCount,
    toggleDrawer,
    clientBlock.fields,
    // DO NOT ADD FORMDATA HERE! Adding formData will kick you out of sub block editors while writing.
--- a/packages/richtext-lexical/src/features/blocks/client/componentInline/index.tsx
+++ b/packages/richtext-lexical/src/features/blocks/client/componentInline/index.tsx
@@ -1,6 +1,6 @@
 'use client'

-import React, { createContext, useCallback, useEffect, useMemo, useRef, useState } from 'react'
+import React, { createContext, useCallback, useEffect, useMemo, useRef } from 'react'
 const baseClass = 'inline-block'

 import type { BlocksFieldClient, FormState } from 'payload'
@@ -70,12 +70,22 @@ export const InlineBlockComponent: React.FC<Props> = (props) => {
  const [editor] = useLexicalComposerContext()
  const { i18n, t } = useTranslation<object, string>()
  const {
-    fieldProps: { featureClientSchemaMap, permissions, readOnly, schemaPath },
+    fieldProps: {
+      featureClientSchemaMap,
+      initialLexicalFormState,
+      permissions,
+      readOnly,
+      schemaPath,
+    },
    uuid: uuidFromContext,
  } = useEditorConfigContext()
  const { getFormState } = useServerFunctions()
  const editDepth = useEditDepth()

+  const [initialState, setInitialState] = React.useState<false | FormState | undefined>(
+    initialLexicalFormState?.[formData.id]?.formState,
+  )
+
  const drawerSlug = formatDrawerSlug({
    slug: `lexical-inlineBlocks-create-` + uuidFromContext,
    depth: editDepth,
@@ -84,7 +94,7 @@ export const InlineBlockComponent: React.FC<Props> = (props) => {

  const inlineBlockElemElemRef = useRef<HTMLDivElement | null>(null)
  const [isSelected, setSelected, clearSelection] = useLexicalNodeSelection(nodeKey)
-  const { id, collectionSlug, docPermissions, getDocPreferences, globalSlug } = useDocumentInfo()
+  const { id, collectionSlug, getDocPreferences, globalSlug } = useDocumentInfo()

  const componentMapRenderedBlockPath = `${schemaPath}.lexical_internal_feature.blocks.lexical_inline_blocks.${formData.blockType}`

@@ -158,46 +168,51 @@ export const InlineBlockComponent: React.FC<Props> = (props) => {

  const onChangeAbortControllerRef = useRef(new AbortController())
  const schemaFieldsPath = `${schemaPath}.lexical_internal_feature.blocks.lexical_inline_blocks.${clientBlock?.slug}.fields`
-  const [initialState, setInitialState] = useState<false | FormState | undefined>(false)

-  /**
-   * HANDLE INITIAL STATE
-   */
+  // Initial state for newly created blocks
  useEffect(() => {
-    const controller = new AbortController()
+    const abortController = new AbortController()

    const awaitInitialState = async () => {
+      /*
+       * This will only run if a new block is created. For all existing blocks that are loaded when the document is loaded, or when the form is saved,
+       * this is not run, as the lexical field RSC will fetch the state server-side and pass it to the client. That way, we avoid unnecessary client-side
+       * requests. Though for newly created blocks, we need to fetch the state client-side, as the server doesn't know about the block yet.
+       */
      const { state } = await getFormState({
        id,
        collectionSlug,
        data: formData,
-        docPermissions,
+        docPermissions: { fields: true },
        docPreferences: await getDocPreferences(),
        globalSlug,
        operation: 'update',
        renderAllFields: true,
        schemaPath: schemaFieldsPath,
-        signal: controller.signal,
+        signal: abortController.signal,
      })

-      setInitialState(state)
+      if (state) {
+        setInitialState(state)
+      }
    }

-    void awaitInitialState()
+    if (formData && !initialState) {
+      void awaitInitialState()
+    }

    return () => {
-      abortAndIgnore(controller)
+      abortAndIgnore(abortController)
    }
  }, [
+    getFormState,
    schemaFieldsPath,
    id,
    formData,
-    getFormState,
+    initialState,
    collectionSlug,
    globalSlug,
-    docPermissions,
    getDocPreferences,
-    clientBlock.slug,
  ])

  /**
@@ -213,7 +228,9 @@ export const InlineBlockComponent: React.FC<Props> = (props) => {
      const { state } = await getFormState({
        id,
        collectionSlug,
-        docPermissions,
+        docPermissions: {
+          fields: true,
+        },
        docPreferences: await getDocPreferences(),
        formState: prevFormState,
        globalSlug,
@@ -228,15 +245,7 @@ export const InlineBlockComponent: React.FC<Props> = (props) => {

      return state
    },
-    [
-      getFormState,
-      id,
-      collectionSlug,
-      docPermissions,
-      getDocPreferences,
-      globalSlug,
-      schemaFieldsPath,
-    ],
+    [getFormState, id, collectionSlug, getDocPreferences, globalSlug, schemaFieldsPath],
  )
  // cleanup effect
  useEffect(() => {
--- a/packages/richtext-lexical/src/features/blocks/client/nodes/BlocksNode.tsx
+++ b/packages/richtext-lexical/src/features/blocks/client/nodes/BlocksNode.tsx
@@ -4,19 +4,10 @@ import type { EditorConfig, LexicalEditor, LexicalNode } from 'lexical'
 import ObjectID from 'bson-objectid'
 import React, { type JSX } from 'react'

-import type {
-  BlockFields,
-  BlockFieldsOptionalID,
-  SerializedBlockNode,
-} from '../../server/nodes/BlocksNode.js'
+import type { BlockFieldsOptionalID, SerializedBlockNode } from '../../server/nodes/BlocksNode.js'

 import { ServerBlockNode } from '../../server/nodes/BlocksNode.js'
-
-const BlockComponent = React.lazy(() =>
-  import('../component/index.js').then((module) => ({
-    default: module.BlockComponent,
-  })),
-)
+import { BlockComponent } from '../component/index.js'

 export class BlockNode extends ServerBlockNode {
  static clone(node: ServerBlockNode): ServerBlockNode {
--- a/packages/richtext-lexical/src/features/typesClient.ts
+++ b/packages/richtext-lexical/src/features/typesClient.ts
@@ -49,18 +49,6 @@ export type PluginComponentWithAnchor<ClientFeatureProps = any> = React.FC<{
 }>

 export type ClientFeature<ClientFeatureProps> = {
-  hooks?: {
-    load?: ({
-      incomingEditorState,
-    }: {
-      incomingEditorState: SerializedEditorState
-    }) => SerializedEditorState
-    save?: ({
-      incomingEditorState,
-    }: {
-      incomingEditorState: SerializedEditorState
-    }) => SerializedEditorState
-  }
  markdownTransformers?: Transformer[]
  nodes?: Array<Klass<LexicalNode> | LexicalNodeReplacement>
  /**
@@ -210,22 +198,6 @@ export type SanitizedPlugin =
 export type SanitizedClientFeatures = {
  /** The keys of all enabled features */
  enabledFeatures: string[]
-  hooks: {
-    load: Array<
-      ({
-        incomingEditorState,
-      }: {
-        incomingEditorState: SerializedEditorState
-      }) => SerializedEditorState
-    >
-    save: Array<
-      ({
-        incomingEditorState,
-      }: {
-        incomingEditorState: SerializedEditorState
-      }) => SerializedEditorState
-    >
-  }
  /**
   * Plugins are react components which get added to the editor. You can use them to interact with lexical, e.g. to create a command which creates a node, or opens a modal, or some other more "outside" functionality
   */
--- a/packages/richtext-lexical/src/field/Field.tsx
+++ b/packages/richtext-lexical/src/field/Field.tsx
@@ -1,5 +1,5 @@
 'use client'
-import type { SerializedEditorState } from 'lexical'
+import type { EditorState, SerializedEditorState } from 'lexical'

 import { FieldLabel, useEditDepth, useField, withCondition } from '@payloadcms/ui'
 import React, { useCallback } from 'react'
@@ -42,7 +42,7 @@ const RichTextComponent: React.FC<
  const memoizedValidate = useCallback(
    (value, validationOptions) => {
      if (typeof validate === 'function') {
-        return validate(value, { ...validationOptions, props, required })
+        return validate(value, { ...validationOptions, required })
      }
    },
    // Important: do not add props to the dependencies array.
@@ -80,6 +80,13 @@ const RichTextComponent: React.FC<

  const pathWithEditDepth = `${path}.${editDepth}`

+  const handleChange = useCallback(
+    (editorState: EditorState) => {
+      setValue(editorState.toJSON())
+    },
+    [setValue],
+  )
+
  return (
    <div
      className={classes}
@@ -99,18 +106,7 @@ const RichTextComponent: React.FC<
            editorConfig={editorConfig}
            fieldProps={props}
            key={JSON.stringify({ initialValue, path })} // makes sure lexical is completely re-rendered when initialValue changes, bypassing the lexical-internal value memoization. That way, external changes to the form will update the editor. More infos in PR description (https://github.com/payloadcms/payload/pull/5010)
-            onChange={(editorState) => {
-              let serializedEditorState = editorState.toJSON()
-
-              // Transform state through save hooks
-              if (editorConfig?.features?.hooks?.save?.length) {
-                editorConfig.features.hooks.save.forEach((hook) => {
-                  serializedEditorState = hook({ incomingEditorState: serializedEditorState })
-                })
-              }
-
-              setValue(serializedEditorState)
-            }}
+            onChange={handleChange}
            readOnly={disabled}
            value={value}
          />
--- a/packages/richtext-lexical/src/field/rscEntry.tsx
+++ b/packages/richtext-lexical/src/field/rscEntry.tsx
@@ -1,10 +1,13 @@
+import type { SerializedLexicalNode } from 'lexical'
 import type {
  ClientComponentProps,
  FieldPaths,
  RichTextFieldClient,
+  RichTextField as RichTextFieldType,
  ServerComponentProps,
 } from 'payload'

+import { renderField } from '@payloadcms/ui/forms/renderField'
 import React from 'react'

 import type { SanitizedServerEditorConfig } from '../lexical/config/types.js'
@@ -12,7 +15,9 @@ import type { LexicalFieldAdminProps } from '../types.js'

 // eslint-disable-next-line payload/no-imports-from-exports-dir
 import { RichTextField } from '../exports/client/index.js'
+import { buildInitialState } from '../utilities/buildInitialState.js'
 import { initLexicalFeatures } from '../utilities/initLexicalFeatures.js'
+
 export const RscEntryLexicalField: React.FC<
  {
    admin: LexicalFieldAdminProps
@@ -20,7 +25,8 @@ export const RscEntryLexicalField: React.FC<
  } & ClientComponentProps &
    Pick<FieldPaths, 'path'> &
    ServerComponentProps
-> = (args) => {
+> = async (args) => {
+  const field: RichTextFieldType = args.field as RichTextFieldType
  const path = args.path ?? (args.clientField as RichTextFieldClient).name
  const schemaPath = args.schemaPath ?? path
  const { clientFeatures, featureClientSchemaMap } = initLexicalFeatures({
@@ -31,6 +37,26 @@ export const RscEntryLexicalField: React.FC<
    sanitizedEditorConfig: args.sanitizedEditorConfig,
    schemaPath,
  })
+
+  let initialLexicalFormState = {}
+  if (args.data?.[field.name]?.root?.children?.length) {
+    initialLexicalFormState = await buildInitialState({
+      context: {
+        id: args.id,
+        collectionSlug: args.collectionSlug,
+        field,
+        fieldSchemaMap: args.fieldSchemaMap,
+        lexicalFieldSchemaPath: schemaPath,
+        operation: args.operation,
+        permissions: args.permissions,
+        preferences: args.preferences,
+        renderFieldFn: renderField,
+        req: args.req,
+      },
+      nodeData: args.data?.[field.name]?.root?.children as SerializedLexicalNode[],
+    })
+  }
+
  return (
    <RichTextField
      admin={args.admin}
@@ -38,6 +64,7 @@ export const RscEntryLexicalField: React.FC<
      featureClientSchemaMap={featureClientSchemaMap}
      field={args.clientField as RichTextFieldClient}
      forceRender={args.forceRender}
+      initialLexicalFormState={initialLexicalFormState}
      lexicalEditorConfig={args.sanitizedEditorConfig.lexical}
      path={path}
      permissions={args.permissions}
--- a/packages/richtext-lexical/src/lexical/LexicalProvider.tsx
+++ b/packages/richtext-lexical/src/lexical/LexicalProvider.tsx
@@ -48,32 +48,22 @@ export const LexicalProvider: React.FC<LexicalProviderProps> = (props) => {

  const editorContainerRef = React.useRef<HTMLDivElement>(null)

-  const processedValue = useMemo(() => {
-    let processed = value
-    if (editorConfig?.features?.hooks?.load?.length) {
-      editorConfig.features.hooks.load.forEach((hook) => {
-        processed = hook({ incomingEditorState: processed })
-      })
-    }
-    return processed
-  }, [editorConfig, value])
-
-  // useMemo for the initialConfig that depends on readOnly and processedValue
+  // useMemo for the initialConfig that depends on readOnly and value
  const initialConfig = useMemo<InitialConfigType>(() => {
-    if (processedValue && typeof processedValue !== 'object') {
+    if (value && typeof value !== 'object') {
      throw new Error(
        'The value passed to the Lexical editor is not an object. This is not supported. Please remove the data from the field and start again. This is the value that was passed in: ' +
-          JSON.stringify(processedValue),
+          JSON.stringify(value),
      )
    }

-    if (processedValue && Array.isArray(processedValue) && !('root' in processedValue)) {
+    if (value && Array.isArray(value) && !('root' in value)) {
      throw new Error(
        'You have tried to pass in data from the old Slate editor to the new Lexical editor. The data structure is different, thus you will have to migrate your data. We offer a one-line migration script which migrates all your rich text fields: https://payloadcms.com/docs/beta/lexical/migration#migration-via-migration-script-recommended',
      )
    }

-    if (processedValue && 'jsonContent' in processedValue) {
+    if (value && 'jsonContent' in value) {
      throw new Error(
        'You have tried to pass in data from payload-plugin-lexical. The data structure is different, thus you will have to migrate your data. Migration guide: https://payloadcms.com/docs/beta/lexical/migration#migrating-from-payload-plugin-lexical',
      )
@@ -81,7 +71,7 @@ export const LexicalProvider: React.FC<LexicalProviderProps> = (props) => {

    return {
      editable: readOnly !== true,
-      editorState: processedValue != null ? JSON.stringify(processedValue) : undefined,
+      editorState: value != null ? JSON.stringify(value) : undefined,
      namespace: editorConfig.lexical.namespace,
      nodes: getEnabledNodes({ editorConfig }),
      onError: (error: Error) => {
@@ -89,7 +79,9 @@ export const LexicalProvider: React.FC<LexicalProviderProps> = (props) => {
      },
      theme: editorConfig.lexical.theme,
    }
-  }, [editorConfig, processedValue, readOnly])
+    // Important: do not add readOnly and value to the dependencies array. This will cause the entire lexical editor to re-render if the document is saved, which will
+    // cause the editor to lose focus.
+  }, [editorConfig])

  if (!initialConfig) {
    return <p>Loading...</p>
--- a/packages/richtext-lexical/src/lexical/config/client/sanitize.ts
+++ b/packages/richtext-lexical/src/lexical/config/client/sanitize.ts
@@ -14,10 +14,6 @@ export const sanitizeClientFeatures = (
 ): SanitizedClientFeatures => {
  const sanitized: SanitizedClientFeatures = {
    enabledFeatures: [],
-    hooks: {
-      load: [],
-      save: [],
-    },
    markdownTransformers: [],
    nodes: [],
    plugins: [],
@@ -39,15 +35,6 @@ export const sanitizeClientFeatures = (
  }

  features.forEach((feature) => {
-    if (feature.hooks) {
-      if (feature.hooks?.load?.length) {
-        sanitized.hooks.load = sanitized.hooks.load.concat(feature.hooks.load)
-      }
-      if (feature.hooks?.save?.length) {
-        sanitized.hooks.save = sanitized.hooks.save.concat(feature.hooks.save)
-      }
-    }
-
    if (feature.providers?.length) {
      sanitized.providers = sanitized.providers.concat(feature.providers)
    }
--- a/packages/richtext-lexical/src/types.ts
+++ b/packages/richtext-lexical/src/types.ts
@@ -15,6 +15,7 @@ import type {
 } from './features/typesClient.js'
 import type { FeatureProviderServer } from './features/typesServer.js'
 import type { SanitizedServerEditorConfig } from './lexical/config/types.js'
+import type { InitialLexicalFormState } from './utilities/buildInitialState.js'

 export type LexicalFieldAdminProps = {
  /**
@@ -95,6 +96,7 @@ export type LexicalRichTextFieldProps = {
    }
  }
  featureClientSchemaMap: FeatureClientSchemaMap
+  initialLexicalFormState: InitialLexicalFormState
  lexicalEditorConfig: LexicalEditorConfig
 } & Pick<ServerFieldBase, 'permissions'> &
  RichTextFieldClientProps<SerializedEditorState, AdapterProps, object>
--- a/packages/richtext-lexical/src/utilities/buildInitialState.ts
+++ b/packages/richtext-lexical/src/utilities/buildInitialState.ts
@@ -0,0 +1,100 @@
+import type { SerializedLexicalNode } from 'lexical'
+import type {
+  DocumentPreferences,
+  FieldSchemaMap,
+  FormState,
+  Operation,
+  PayloadRequest,
+  RichTextField,
+  SanitizedFieldPermissions,
+} from 'payload'
+
+import { fieldSchemasToFormState } from '@payloadcms/ui/forms/fieldSchemasToFormState'
+
+import type { SerializedBlockNode } from '../nodeTypes.js'
+
+export type InitialLexicalFormState = {
+  [nodeID: string]: {
+    [key: string]: any
+    formState?: FormState
+  }
+}
+
+type Props = {
+  context: {
+    collectionSlug: string
+    field: RichTextField
+    fieldSchemaMap: FieldSchemaMap
+    id?: number | string
+    lexicalFieldSchemaPath: string
+    operation: Operation
+    permissions?: SanitizedFieldPermissions
+    preferences: DocumentPreferences
+    renderFieldFn: any
+    req: PayloadRequest
+  }
+  initialState?: InitialLexicalFormState
+  nodeData: SerializedLexicalNode[]
+}
+
+export async function buildInitialState({
+  context,
+  initialState: initialStateFromArgs,
+  nodeData,
+}: Props): Promise<InitialLexicalFormState> {
+  let initialState: InitialLexicalFormState = initialStateFromArgs ?? {}
+  for (const node of nodeData) {
+    if ('children' in node) {
+      initialState = await buildInitialState({
+        context,
+        initialState,
+        nodeData: node.children as SerializedLexicalNode[],
+      })
+    }
+
+    if (node.type === 'block' || node.type === 'inlineBlock') {
+      const blockNode = node as SerializedBlockNode
+      const id = blockNode?.fields?.id
+      if (!id) {
+        continue
+      }
+
+      const schemaFieldsPath =
+        node.type === 'block'
+          ? `${context.lexicalFieldSchemaPath}.lexical_internal_feature.blocks.lexical_blocks.${blockNode.fields.blockType}.fields`
+          : `${context.lexicalFieldSchemaPath}.lexical_internal_feature.blocks.lexical_inline_blocks.${blockNode.fields.blockType}.fields`
+
+      // Build form state for the block
+
+      const formStateResult = await fieldSchemasToFormState({
+        id: context.id,
+        collectionSlug: context.collectionSlug,
+        data: blockNode.fields,
+        fields: (context.fieldSchemaMap.get(schemaFieldsPath) as any)?.fields,
+        fieldSchemaMap: context.fieldSchemaMap,
+        operation: context.operation as any, // TODO: Type
+        permissions: true,
+        preferences: context.preferences,
+        renderAllFields: true, // If this function runs, the parent lexical field is being re-rendered => thus we can assume all its sub-fields need to be re-rendered
+        renderFieldFn: context.renderFieldFn,
+        req: context.req,
+        schemaPath: schemaFieldsPath,
+      })
+
+      if (!initialState[id]) {
+        initialState[id] = {}
+      }
+
+      initialState[id].formState = formStateResult
+
+      if (node.type === 'block') {
+        const currentFieldPreferences = context.preferences?.fields[context.field.name]
+        const collapsedArray = currentFieldPreferences?.collapsed
+        if (collapsedArray && collapsedArray.includes(id)) {
+          initialState[id].collapsed = true
+        }
+      }
+    }
+  }
+  return initialState
+}
--- a/packages/richtext-lexical/src/utilities/fieldsDrawer/DrawerContent.tsx
+++ b/packages/richtext-lexical/src/utilities/fieldsDrawer/DrawerContent.tsx
@@ -27,7 +27,7 @@ export const DrawerContent: React.FC<Omit<FieldsDrawerProps, 'drawerSlug' | 'dra
  schemaPathSuffix,
 }) => {
  const { t } = useTranslation()
-  const { id, collectionSlug, docPermissions, getDocPreferences, globalSlug } = useDocumentInfo()
+  const { id, collectionSlug, getDocPreferences, globalSlug } = useDocumentInfo()

  const onChangeAbortControllerRef = useRef(new AbortController())

@@ -53,7 +53,9 @@ export const DrawerContent: React.FC<Omit<FieldsDrawerProps, 'drawerSlug' | 'dra
        id,
        collectionSlug,
        data: data ?? {},
-        docPermissions,
+        docPermissions: {
+          fields: true,
+        },
        docPreferences: await getDocPreferences(),
        globalSlug,
        operation: 'update',
@@ -70,16 +72,7 @@ export const DrawerContent: React.FC<Omit<FieldsDrawerProps, 'drawerSlug' | 'dra
    return () => {
      abortAndIgnore(controller)
    }
-  }, [
-    schemaFieldsPath,
-    id,
-    data,
-    getFormState,
-    collectionSlug,
-    globalSlug,
-    docPermissions,
-    getDocPreferences,
-  ])
+  }, [schemaFieldsPath, id, data, getFormState, collectionSlug, globalSlug, getDocPreferences])

  const onChange = useCallback(
    async ({ formState: prevFormState }) => {
@@ -91,7 +84,9 @@ export const DrawerContent: React.FC<Omit<FieldsDrawerProps, 'drawerSlug' | 'dra
      const { state } = await getFormState({
        id,
        collectionSlug,
-        docPermissions,
+        docPermissions: {
+          fields: true,
+        },
        docPreferences: await getDocPreferences(),
        formState: prevFormState,
        globalSlug,
@@ -106,15 +101,7 @@ export const DrawerContent: React.FC<Omit<FieldsDrawerProps, 'drawerSlug' | 'dra

      return state
    },
-    [
-      getFormState,
-      id,
-      collectionSlug,
-      docPermissions,
-      getDocPreferences,
-      globalSlug,
-      schemaFieldsPath,
-    ],
+    [getFormState, id, collectionSlug, getDocPreferences, globalSlug, schemaFieldsPath],
  )

  // cleanup effect
--- a/packages/ui/package.json
+++ b/packages/ui/package.json
@@ -63,6 +63,11 @@
      "types": "./src/forms/fieldSchemasToFormState/index.tsx",
      "default": "./src/forms/fieldSchemasToFormState/index.tsx"
    },
+    "./forms/renderField": {
+      "import": "./src/forms/fieldSchemasToFormState/renderField.tsx",
+      "types": "./src/forms/fieldSchemasToFormState/renderField.tsx",
+      "default": "./src/forms/fieldSchemasToFormState/renderField.tsx"
+    },
    "./scss/app.scss": "./src/scss/app.scss",
    "./assets": {
      "import": "./src/assets/index.ts",
@@ -190,6 +195,16 @@
        "types": "./dist/fields/*/index.d.ts",
        "default": "./dist/fields/*/index.js"
      },
+      "./forms/fieldSchemasToFormState": {
+        "import": "./dist/forms/fieldSchemasToFormState/index.js",
+        "types": "./dist/forms/fieldSchemasToFormState/index.d.ts",
+        "default": "./dist/forms/fieldSchemasToFormState/index.js"
+      },
+      "./forms/renderField": {
+        "import": "./dist/forms/fieldSchemasToFormState/renderField.js",
+        "types": "./dist/forms/fieldSchemasToFormState/renderField.d.ts",
+        "default": "./dist/forms/fieldSchemasToFormState/renderField.js"
+      },
      "./forms/*": {
        "import": "./dist/forms/*/index.js",
        "types": "./dist/forms/*/index.d.ts",
--- a/packages/ui/src/forms/fieldSchemasToFormState/addFieldStatePromise.ts
+++ b/packages/ui/src/forms/fieldSchemasToFormState/addFieldStatePromise.ts
@@ -715,6 +715,8 @@ export const addFieldStatePromise = async (args: AddFieldStatePromiseArgs): Prom
    }

    renderFieldFn({
+      id,
+      collectionSlug,
      data: fullData,
      fieldConfig: fieldConfig as Field,
      fieldSchemaMap,
@@ -726,6 +728,7 @@ export const addFieldStatePromise = async (args: AddFieldStatePromiseArgs): Prom
      parentSchemaPath,
      path,
      permissions,
+      preferences,
      previousFieldState: previousFormState?.[path],
      req,
      schemaPath,
--- a/packages/ui/src/forms/fieldSchemasToFormState/renderField.tsx
+++ b/packages/ui/src/forms/fieldSchemasToFormState/renderField.tsx
@@ -14,7 +14,8 @@ import { fieldAffectsData } from 'payload/shared'
 import type { RenderFieldMethod } from './types.js'

 import { RenderServerComponent } from '../../elements/RenderServerComponent/index.js'
-import { FieldDescription } from '../../fields/FieldDescription/index.js'
+// eslint-disable-next-line payload/no-imports-from-exports-dir -- need this to reference already existing bundle. Otherwise, bundle size increases., payload/no-imports-from-exports-dir
+import { FieldDescription } from '../../exports/client/index.js'

 const defaultUIFieldComponentKeys: Array<'Cell' | 'Description' | 'Field' | 'Filter'> = [
  'Cell',
@@ -24,6 +25,8 @@ const defaultUIFieldComponentKeys: Array<'Cell' | 'Description' | 'Field' | 'Fil
 ]

 export const renderField: RenderFieldMethod = ({
+  id,
+  collectionSlug,
  data,
  fieldConfig,
  fieldSchemaMap,
@@ -35,6 +38,7 @@ export const renderField: RenderFieldMethod = ({
  parentSchemaPath,
  path,
  permissions: incomingPermissions,
+  preferences,
  req,
  schemaPath,
  siblingData,
@@ -72,6 +76,7 @@ export const renderField: RenderFieldMethod = ({
  }

  const serverProps: ServerComponentProps = {
+    id,
    clientField,
    data,
    field: fieldConfig,
@@ -79,9 +84,13 @@ export const renderField: RenderFieldMethod = ({
    permissions,
    // TODO: Should we pass explicit values? initialValue, value, valid
    // value and initialValue should be typed
+    collectionSlug,
    formState,
    i18n: req.i18n,
+    operation,
    payload: req.payload,
+    preferences,
+    req,
    siblingData,
    user: req.user,
  }
--- a/packages/ui/src/forms/fieldSchemasToFormState/types.ts
+++ b/packages/ui/src/forms/fieldSchemasToFormState/types.ts
@@ -1,5 +1,6 @@
 import type {
  Data,
+  DocumentPreferences,
  Field,
  FieldSchemaMap,
  FieldState,
@@ -10,11 +11,13 @@ import type {
 } from 'payload'

 export type RenderFieldArgs = {
+  collectionSlug: string
  data: Data
  fieldConfig: Field
  fieldSchemaMap: FieldSchemaMap
  fieldState: FieldState
  formState: FormState
+  id?: number | string
  indexPath: string
  operation: Operation
  parentPath: string
@@ -26,6 +29,7 @@ export type RenderFieldArgs = {
      }
    | null
    | SanitizedFieldPermissions
+  preferences: DocumentPreferences
  previousFieldState: FieldState
  req: PayloadRequest
  schemaPath: string
--- a/test/fields/collections/Lexical/e2e/main/e2e.spec.ts
+++ b/test/fields/collections/Lexical/e2e/main/e2e.spec.ts
@@ -14,6 +14,7 @@ import {
  ensureCompilationIsDone,
  initPageConsoleErrorCatch,
  saveDocAndAssert,
+  saveDocHotkeyAndAssert,
  throttleTest,
 } from '../../../../../helpers.js'
 import { AdminUrlUtil } from '../../../../../helpers/adminUrlUtil.js'
@@ -201,6 +202,63 @@ describe('lexicalMain', () => {
      timeout: POLL_TOPASS_TIMEOUT,
    })
  })
+
+  test('ensure saving document does not kick cursor / focus out of rich text field', async () => {
+    await navigateToLexicalFields()
+    const richTextField = page.locator('.rich-text-lexical').nth(2) // second
+    await richTextField.scrollIntoViewIfNeeded()
+    await expect(richTextField).toBeVisible()
+    // Wait until there at least 10 blocks visible in that richtext field - thus wait for it to be fully loaded
+    await expect(richTextField.locator('.lexical-block')).toHaveCount(10)
+    await expect(page.locator('.shimmer-effect')).toHaveCount(0)
+
+    const spanInEditor = richTextField.locator('span').getByText('Upload Node:').first()
+    await expect(spanInEditor).toBeVisible()
+
+    await spanInEditor.click() // Click works better than focus
+    // Now go to the END of the span
+    for (let i = 0; i < 6; i++) {
+      await page.keyboard.press('ArrowRight')
+    }
+
+    await page.keyboard.type('more')
+    await expect(spanInEditor).toHaveText('Upload Node:more')
+
+    await wait(500)
+
+    await saveDocHotkeyAndAssert(page) // Use hotkey to save, as clicking the save button will obviously remove focus from the richtext field
+    await wait(500)
+    // Keep writing after save, assuming the cursor position is still at the end of the span
+    await page.keyboard.type('text')
+    await expect(spanInEditor).toHaveText('Upload Node:moretext')
+    await wait(500)
+    await saveDocAndAssert(page) // Use hotkey to save, as clicking the save button will obviously remove focus from the richtext field
+
+    await expect(async () => {
+      const lexicalDoc: LexicalField = (
+        await payload.find({
+          collection: lexicalFieldsSlug,
+          depth: 0,
+          overrideAccess: true,
+          where: {
+            title: {
+              equals: lexicalDocData.title,
+            },
+          },
+        })
+      ).docs[0] as never
+
+      const lexicalField: SerializedEditorState = lexicalDoc.lexicalWithBlocks
+      const firstParagraphTextNode: SerializedTextNode = (
+        lexicalField.root.children[0] as SerializedParagraphNode
+      ).children[0] as SerializedTextNode
+
+      expect(firstParagraphTextNode.text).toBe('Upload Node:moretext')
+    }).toPass({
+      timeout: POLL_TOPASS_TIMEOUT,
+    })
+  })
+
  test('should be able to bold text using floating select toolbar', async () => {
    await navigateToLexicalFields()
    const richTextField = page.locator('.rich-text-lexical').nth(2) // second
--- a/test/helpers.ts
+++ b/test/helpers.ts
@@ -209,7 +209,12 @@ export async function saveDocHotkeyAndAssert(page: Page): Promise<void> {
  } else {
    await page.keyboard.down('Control')
  }
-  await page.keyboard.down('s')
+  await page.keyboard.press('s')
+  if (isMac) {
+    await page.keyboard.up('Meta')
+  } else {
+    await page.keyboard.up('Control')
+  }
  await expect(page.locator('.payload-toast-container')).toContainText('successfully')
 }