fix(richtext-*): fix client features were not loaded properly, improve performance of LexicalProvider, slate cell component was non-functional, support richtext adapter Cell RSCs (#6573)

packages/payload/src/admin/elements/Cell.ts

@@ -12,13 +12,6 @@ import type {
 export type RowData = Record<string, any>
 
 export type CellComponentProps = {
-  /**
-   * A custom component to override the default cell component. If this is not set, the React component will be
-   * taken from cellComponents based on the field type.
-   *
-   * This is used to provide the RichText cell component for the RichText field.
-   */
-  CellComponentOverride?: React.ReactNode
   blocks?: {
     labels: BlockField['labels']
     slug: string
@@ -39,6 +32,7 @@ export type CellComponentProps = {
   options?: SelectField['options']
   relationTo?: RelationshipField['relationTo']
   richTextComponentMap?: Map<string, React.ReactNode> // any should be MappedField
+  schemaPath: string
 }
 
 export type DefaultCellComponentProps<T = any> = CellComponentProps & {

packages/richtext-lexical/src/cell/index.tsx

@@ -1,9 +1,9 @@
 'use client'
 import type { EditorConfig as LexicalEditorConfig } from 'lexical'
+import type { CellComponentProps } from 'payload/types'
 
 import { createHeadlessEditor } from '@lexical/headless'
 import { useTableCell } from '@payloadcms/ui/elements/Table'
-import { useFieldProps } from '@payloadcms/ui/forms/FieldPropsProvider'
 import { useClientFunctions } from '@payloadcms/ui/providers/ClientFunction'
 import { $getRoot } from 'lexical'
 import React, { useEffect, useState } from 'react'
@@ -17,16 +17,20 @@ import { loadClientFeatures } from '../field/lexical/config/client/loader.js'
 import { sanitizeClientEditorConfig } from '../field/lexical/config/client/sanitize.js'
 import { getEnabledNodes } from '../field/lexical/nodes/index.js'
 
-export const RichTextCell: React.FC<{
-  admin?: LexicalFieldAdminProps
-  lexicalEditorConfig: LexicalEditorConfig
-}> = (props) => {
-  const { admin, lexicalEditorConfig } = props
+export const RichTextCell: React.FC<
+  CellComponentProps & {
+    admin?: LexicalFieldAdminProps
+    lexicalEditorConfig: LexicalEditorConfig
+  }
+> = (props) => {
+  const { admin, lexicalEditorConfig, richTextComponentMap } = props
 
   const [preview, setPreview] = React.useState('Loading...')
-  const { schemaPath } = useFieldProps()
 
-  const { cellData, richTextComponentMap } = useTableCell()
+  const {
+    cellData,
+    cellProps: { schemaPath },
+  } = useTableCell()
 
   const clientFunctions = useClientFunctions()
   const [hasLoadedFeatures, setHasLoadedFeatures] = useState(false)
@@ -36,21 +40,24 @@ export const RichTextCell: React.FC<{
   const [finalSanitizedEditorConfig, setFinalSanitizedEditorConfig] =
     useState<SanitizedClientEditorConfig>(null)
 
-  let featureProviderComponents: GeneratedFeatureProviderComponent[] = richTextComponentMap.get(
-    'features',
-  ) as GeneratedFeatureProviderComponent[] // TODO: Type better
-  // order by order
-  featureProviderComponents = featureProviderComponents.sort((a, b) => a.order - b.order)
+  const featureProviderComponents: GeneratedFeatureProviderComponent[] = (
+    richTextComponentMap.get('features') as GeneratedFeatureProviderComponent[]
+  ).sort((a, b) => a.order - b.order) // order by order
 
-  const featureComponentsWithFeaturesLength =
-    Array.from(richTextComponentMap.keys()).filter(
-      (key) => key.startsWith(`feature.`) && !key.includes('.fields.'),
-    ).length + featureProviderComponents.length
+  let featureProvidersAndComponentsToLoad = 0 // feature providers and components
+  for (const featureProvider of featureProviderComponents) {
+    const featureComponentKeys = Array.from(richTextComponentMap.keys()).filter((key) =>
+      key.startsWith(`feature.${featureProvider.key}.components.`),
+    )
+
+    featureProvidersAndComponentsToLoad += 1
+    featureProvidersAndComponentsToLoad += featureComponentKeys.length
+  }
 
   useEffect(() => {
     if (!hasLoadedFeatures) {
       const featureProvidersLocal: FeatureProviderClient<unknown>[] = []
-      let featureProvidersAndComponentsLoaded = 0
+      let featureProvidersAndComponentsLoaded = 0 // feature providers and components only
 
       Object.entries(clientFunctions).forEach(([key, plugin]) => {
         if (key.startsWith(`lexicalFeature.${schemaPath}.`)) {
@@ -61,7 +68,7 @@ export const RichTextCell: React.FC<{
         }
       })
 
-      if (featureProvidersAndComponentsLoaded === featureComponentsWithFeaturesLength) {
+      if (featureProvidersAndComponentsLoaded === featureProvidersAndComponentsToLoad) {
         setFeatureProviders(featureProvidersLocal)
         setHasLoadedFeatures(true)
 
@@ -89,6 +96,7 @@ export const RichTextCell: React.FC<{
     }
   }, [
     admin,
+    featureProviderComponents,
     hasLoadedFeatures,
     clientFunctions,
     schemaPath,
@@ -96,10 +104,14 @@ export const RichTextCell: React.FC<{
     featureProviders,
     finalSanitizedEditorConfig,
     lexicalEditorConfig,
-    featureComponentsWithFeaturesLength,
+    richTextComponentMap,
+    featureProvidersAndComponentsToLoad,
   ])
 
   useEffect(() => {
+    if (!hasLoadedFeatures) {
+      return
+    }
     let dataToUse = cellData
     if (dataToUse == null || !hasLoadedFeatures || !finalSanitizedEditorConfig) {
       setPreview('')
@@ -156,6 +168,7 @@ export const RichTextCell: React.FC<{
             const featureComponentKeys = Array.from(richTextComponentMap.keys()).filter((key) =>
               key.startsWith(`feature.${featureProvider.key}.components.`),
             )
+
             const featureComponents: React.ReactNode[] = featureComponentKeys.map((key) => {
               return richTextComponentMap.get(key)
             })

packages/richtext-lexical/src/field/features/createFeaturePropComponent.tsx

@@ -2,11 +2,17 @@
 
 import type React from 'react'
 
+import { useTableCell } from '@payloadcms/ui/elements/Table'
 import { useFieldProps } from '@payloadcms/ui/forms/FieldPropsProvider'
 import { useAddClientFunction } from '@payloadcms/ui/providers/ClientFunction'
 
 const useLexicalFeatureProp = <T,>(featureKey: string, componentKey: string, prop: T) => {
-  const { schemaPath } = useFieldProps()
+  const { schemaPath: schemaPathFromFieldProps } = useFieldProps()
+  const tableCell = useTableCell()
+
+  const schemaPathFromCellProps = tableCell?.cellProps?.schemaPath
+
+  const schemaPath = schemaPathFromCellProps || schemaPathFromFieldProps // schemaPathFromCellProps needs to have priority, as there can be cells within fields (e.g. list drawers) and the cell schemaPath needs to be used there - not the parent field schemaPath. There cannot be fields within cells.
 
   useAddClientFunction(
     `lexicalFeature.${schemaPath}.${featureKey}.components.${componentKey}`,

packages/richtext-lexical/src/field/index.tsx

@@ -39,14 +39,19 @@ export const RichTextField: React.FC<
 
   let featureProviderComponents: GeneratedFeatureProviderComponent[] = richTextComponentMap.get(
     'features',
-  ) as GeneratedFeatureProviderComponent[] // TODO: Type better
+  ) as GeneratedFeatureProviderComponent[]
   // order by order
   featureProviderComponents = featureProviderComponents.sort((a, b) => a.order - b.order)
 
-  const featureComponentsWithFeaturesLength =
-    Array.from(richTextComponentMap.keys()).filter(
-      (key) => key.startsWith(`feature.`) && !key.includes('.fields.'),
-    ).length + featureProviderComponents.length
+  let featureProvidersAndComponentsToLoad = 0 // feature providers and components
+  for (const featureProvider of featureProviderComponents) {
+    const featureComponentKeys = Array.from(richTextComponentMap.keys()).filter((key) =>
+      key.startsWith(`feature.${featureProvider.key}.components.`),
+    )
+
+    featureProvidersAndComponentsToLoad += 1
+    featureProvidersAndComponentsToLoad += featureComponentKeys.length
+  }
 
   useEffect(() => {
     if (!hasLoadedFeatures) {
@@ -62,7 +67,7 @@ export const RichTextField: React.FC<
         }
       })
 
-      if (featureProvidersAndComponentsLoaded === featureComponentsWithFeaturesLength) {
+      if (featureProvidersAndComponentsLoaded === featureProvidersAndComponentsToLoad) {
         setFeatureProviders(featureProvidersLocal)
         setHasLoadedFeatures(true)
 
@@ -97,7 +102,7 @@ export const RichTextField: React.FC<
     featureProviders,
     finalSanitizedEditorConfig,
     lexicalEditorConfig,
-    featureComponentsWithFeaturesLength,
+    featureProvidersAndComponentsToLoad,
   ])
 
   if (!hasLoadedFeatures) {

packages/richtext-lexical/src/field/lexical/LexicalProvider.tsx

@@ -6,6 +6,7 @@ import type { LexicalEditor } from 'lexical'
 
 import { LexicalComposer } from '@lexical/react/LexicalComposer.js'
 import * as React from 'react'
+import { useMemo } from 'react'
 
 import type { SanitizedClientEditorConfig } from './config/types.js'
 
@@ -29,19 +30,47 @@ export type LexicalProviderProps = {
   value: SerializedEditorState
 }
 export const LexicalProvider: React.FC<LexicalProviderProps> = (props) => {
-  const { editorConfig, fieldProps, onChange, path, readOnly } = props
-  let { value } = props
+  const { editorConfig, fieldProps, value, onChange, path, readOnly } = props
+
   const parentContext = useEditorConfigContext()
 
   const editorContainerRef = React.useRef<HTMLDivElement>(null)
 
-  const [initialConfig, setInitialConfig] = React.useState<InitialConfigType | null>(null)
 
-  // set lexical config in useEffect: // TODO: Is this the most performant way to do this? Prob not
-  React.useEffect(() => {
-    const newInitialConfig: InitialConfigType = {
+  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
+  const initialConfig = useMemo<InitialConfigType>(() => {
+    if (processedValue && typeof processedValue !== '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),
+      )
+    }
+
+    if (processedValue && Array.isArray(processedValue) && !('root' in processedValue)) {
+      throw new Error(
+        'You have tried to pass in data from the old, Slate editor, to the new, Lexical editor. This is not supported. There is no automatic conversion from Slate to Lexical data available yet (coming soon). Please remove the data from the field and start again.',
+      )
+    }
+
+    if (processedValue && 'jsonContent' in processedValue) {
+      throw new Error(
+        'You have tried to pass in data from payload-plugin-lexical. This is not supported. The data structure has changed in this editor, compared to the plugin, and there is no automatic conversion available yet (coming soon). Please remove the data from the field and start again.',
+      )
+    }
+
+    return {
       editable: readOnly !== true,
-      editorState: value != null ? JSON.stringify(value) : undefined,
+      editorState: processedValue != null ? JSON.stringify(processedValue) : undefined,
       namespace: editorConfig.lexical.namespace,
       nodes: [...getEnabledNodes({ editorConfig })],
       onError: (error: Error) => {
@@ -49,33 +78,7 @@ export const LexicalProvider: React.FC<LexicalProviderProps> = (props) => {
       },
       theme: editorConfig.lexical.theme,
     }
-    setInitialConfig(newInitialConfig)
-  }, [editorConfig, readOnly, value])
-
-  if (editorConfig?.features?.hooks?.load?.length) {
-    editorConfig.features.hooks.load.forEach((hook) => {
-      value = hook({ incomingEditorState: value })
-    })
-  }
-
-  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(value),
-    )
-  }
-
-  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. This is not supported. There is no automatic conversion from Slate to Lexical data available yet (coming soon). Please remove the data from the field and start again.',
-    )
-  }
-
-  if (value && 'jsonContent' in value) {
-    throw new Error(
-      'You have tried to pass in data from payload-plugin-lexical. This is not supported. The data structure has changed in this editor, compared to the plugin, and there is no automatic conversion available yet (coming soon). Please remove the data from the field and start again.',
-    )
-  }
+  }, [editorConfig, processedValue, readOnly])
 
   if (!initialConfig) {
     return <p>Loading...</p>

packages/richtext-lexical/src/useLexicalFeature.tsx

@@ -1,5 +1,6 @@
 'use client'
 
+import { useTableCell } from '@payloadcms/ui/elements/Table'
 import { useFieldProps } from '@payloadcms/ui/forms/FieldPropsProvider'
 import { useAddClientFunction } from '@payloadcms/ui/providers/ClientFunction'
 
@@ -9,7 +10,12 @@ export const useLexicalFeature = <ClientFeatureProps,>(
   featureKey: string,
   feature: FeatureProviderClient<ClientFeatureProps>,
 ) => {
-  const { schemaPath } = useFieldProps()
+  const { schemaPath: schemaPathFromFieldProps } = useFieldProps()
+  const tableCell = useTableCell()
+
+  const schemaPathFromCellProps = tableCell?.cellProps?.schemaPath
+
+  const schemaPath = schemaPathFromCellProps || schemaPathFromFieldProps // schemaPathFromCellProps needs to have priority, as there can be cells within fields (e.g. list drawers) and the cell schemaPath needs to be used there - not the parent field schemaPath. There cannot be fields within cells.
 
   useAddClientFunction(`lexicalFeature.${schemaPath}.${featureKey}`, feature)
 }

packages/richtext-slate/src/cell/index.tsx

@@ -1,9 +1,11 @@
 'use client'
 import type { DefaultCellComponentProps } from 'payload/types'
 
+import { useTableCell } from '@payloadcms/ui/elements/Table'
 import React from 'react'
 
-export const RichTextCell: React.FC<DefaultCellComponentProps<any[]>> = ({ cellData }) => {
+export const RichTextCell: React.FC<DefaultCellComponentProps<any[]>> = () => {
+  const { cellData } = useTableCell()
   const flattenedText = cellData?.map((i) => i?.children?.map((c) => c.text)).join(' ')
 
   // Limiting the number of characters shown is done in a CSS rule

packages/ui/src/elements/Table/DefaultCell/fields/Textarea/index.tsx

@@ -4,6 +4,6 @@ import type { DefaultCellComponentProps } from 'payload/types'
 import React from 'react'
 
 export const TextareaCell: React.FC<DefaultCellComponentProps<string>> = ({ cellData }) => {
-  const textToShow = cellData?.length > 100 ? `${cellData.substr(0, 100)}\u2026` : cellData
+  const textToShow = cellData?.length > 100 ? `${cellData.substring(0, 100)}\u2026` : cellData
   return <span>{textToShow}</span>
 }

packages/ui/src/elements/Table/DefaultCell/index.tsx

@@ -8,7 +8,7 @@ import { getTranslation } from '@payloadcms/translations'
 
 import { useConfig } from '../../../providers/Config/index.js'
 import { useTranslation } from '../../../providers/Translation/index.js'
-import { TableCellProvider, useTableCell } from '../TableCellProvider/index.js'
+import { useTableCell } from '../TableCellProvider/index.js'
 import { CodeCell } from './fields/Code/index.js'
 import { cellComponents } from './fields/index.js'
 
@@ -17,13 +17,11 @@ const Link = (LinkImport.default || LinkImport) as unknown as typeof LinkImport.
 export const DefaultCell: React.FC<CellComponentProps> = (props) => {
   const {
     name,
-    CellComponentOverride,
     className: classNameFromProps,
     fieldType,
     isFieldAffectingData,
     label,
     onClick: onClickFromProps,
-    richTextComponentMap,
   } = props
 
   const { i18n } = useTranslation()
@@ -77,7 +75,13 @@ export const DefaultCell: React.FC<CellComponentProps> = (props) => {
   if (name === 'id') {
     return (
       <WrapElement {...wrapElementProps}>
-        <CodeCell cellData={`ID: ${cellData}`} name={name} nowrap rowData={rowData} />
+        <CodeCell
+          cellData={`ID: ${cellData}`}
+          name={name}
+          nowrap
+          rowData={rowData}
+          schemaPath={cellContext?.cellProps?.schemaPath}
+        />
       </WrapElement>
     )
   }
@@ -85,15 +89,9 @@ export const DefaultCell: React.FC<CellComponentProps> = (props) => {
   const DefaultCellComponent: React.FC<DefaultCellComponentProps> =
     typeof cellData !== 'undefined' && cellComponents[fieldType]
 
-  let CellComponent: React.ReactNode =
-    cellData &&
-    (CellComponentOverride ? ( // CellComponentOverride is used for richText
-      <TableCellProvider richTextComponentMap={richTextComponentMap}>
-        {CellComponentOverride}
-      </TableCellProvider>
-    ) : null)
+  let CellComponent: React.ReactNode = null
 
-  if (!CellComponent && DefaultCellComponent) {
+  if (DefaultCellComponent) {
     CellComponent = (
       <DefaultCellComponent
         cellData={cellData}
@@ -102,7 +100,7 @@ export const DefaultCell: React.FC<CellComponentProps> = (props) => {
         {...props}
       />
     )
-  } else if (!CellComponent && !DefaultCellComponent) {
+  } else if (!DefaultCellComponent) {
     // DefaultCellComponent does not exist for certain field types like `text`
     if (customCellContext.uploadConfig && isFieldAffectingData && name === 'filename') {
       const FileCellComponent = cellComponents.File

packages/ui/src/elements/Table/TableCellProvider/index.tsx

@@ -8,7 +8,6 @@ export type ITableCellContext = {
   cellProps?: Partial<CellComponentProps>
   columnIndex?: number
   customCellContext: DefaultCellComponentProps['customCellContext']
-  richTextComponentMap?: DefaultCellComponentProps['richTextComponentMap']
   rowData: DefaultCellComponentProps['rowData']
 }
 
@@ -20,18 +19,9 @@ export const TableCellProvider: React.FC<{
   children: React.ReactNode
   columnIndex?: number
   customCellContext?: DefaultCellComponentProps['customCellContext']
-  richTextComponentMap?: DefaultCellComponentProps['richTextComponentMap']
   rowData?: DefaultCellComponentProps['rowData']
 }> = (props) => {
-  const {
-    cellData,
-    cellProps,
-    children,
-    columnIndex,
-    customCellContext,
-    richTextComponentMap,
-    rowData,
-  } = props
+  const { cellData, cellProps, children, columnIndex, customCellContext, rowData } = props
 
   const contextToInherit = useTableCell()
 
@@ -44,7 +34,6 @@ export const TableCellProvider: React.FC<{
         customCellContext,
         rowData,
         ...contextToInherit,
-        richTextComponentMap,
       }}
     >
       {children}

packages/ui/src/elements/withMergedProps/index.tsx

@@ -1,5 +1,5 @@
 import { serverProps } from 'payload/config'
-import { deepMerge } from 'payload/utilities'
+import { deepMerge, isReactServerComponentOrFunction } from 'payload/utilities'
 import React from 'react'
 
 /**
@@ -19,23 +19,30 @@ import React from 'react'
  * // <OriginalComponent customProp="value" someExtraValue={5} />
  *
  * @returns A higher-order component with combined properties.
+ *
+ * @param Component - The original component to wrap.
+ * @param sanitizeServerOnlyProps - If true, server-only props will be removed from the merged props. @default true if the component is not a server component, false otherwise.
+ * @param toMergeIntoProps - The properties to merge into the passed props.
  */
 export function withMergedProps<ToMergeIntoProps, CompleteReturnProps>({
   Component,
-  sanitizeServerOnlyProps = true,
+  sanitizeServerOnlyProps,
   toMergeIntoProps,
 }: {
   Component: React.FC<CompleteReturnProps>
   sanitizeServerOnlyProps?: boolean
   toMergeIntoProps: ToMergeIntoProps
 }): React.FC<CompleteReturnProps> {
+  if (sanitizeServerOnlyProps === undefined) {
+    sanitizeServerOnlyProps = !isReactServerComponentOrFunction(Component)
+  }
   // A wrapper around the args.Component to inject the args.toMergeArgs as props, which are merged with the passed props
   const MergedPropsComponent: React.FC<CompleteReturnProps> = (passedProps) => {
     const mergedProps = deepMerge(passedProps, toMergeIntoProps)
 
     if (sanitizeServerOnlyProps) {
       serverProps.forEach((prop) => {
-        delete (mergedProps)[prop]
+        delete mergedProps[prop]
       })
     }
 

packages/ui/src/providers/ComponentMap/buildComponentMap/fields.tsx

@@ -75,7 +75,7 @@ export const mapFields = (args: {
     const fieldIsPresentational = fieldIsPresentationalOnly(field)
     let CustomFieldComponent: CustomComponent<FieldComponentProps> = field.admin?.components?.Field
 
-    const CustomCellComponent = field.admin?.components?.Cell
+    let CustomCellComponent = field.admin?.components?.Cell
 
     const isHidden = field?.admin && 'hidden' in field.admin && field.admin.hidden
 
@@ -238,6 +238,7 @@ export const mapFields = (args: {
           labels: 'labels' in field ? field.labels : undefined,
           options: 'options' in field ? fieldOptions : undefined,
           relationTo: 'relationTo' in field ? field.relationTo : undefined,
+          schemaPath: path,
         }
 
         switch (field.type) {
@@ -588,9 +589,7 @@ export const mapFields = (args: {
             }
 
             if (RichTextCellComponent) {
-              cellComponentProps.CellComponentOverride = (
-                <WithServerSideProps Component={RichTextCellComponent} />
-              )
+              CustomCellComponent = RichTextCellComponent
             }
 
             fieldComponentProps = richTextField
@@ -788,6 +787,7 @@ export const mapFields = (args: {
       CustomField: null,
       cellComponentProps: {
         name: 'id',
+        schemaPath: 'id',
       },
       disableBulkEdit: true,
       fieldComponentProps: {

test/fields/collections/Lexical/e2e/blocks/e2e.spec.ts

@@ -530,33 +530,52 @@ describe('lexicalBlocks', () => {
       await expect(paragraphInSubEditor).toBeVisible()
       await paragraphInSubEditor.click()
       await page.keyboard.type('Some subText')
+
       // Upload something
-      const chooseExistingUploadButton = newSubLexicalAndUploadBlock
-        .locator('.upload__toggler.list-drawer__toggler')
-        .first()
-      await expect(chooseExistingUploadButton).toBeVisible()
-      await chooseExistingUploadButton.click()
-      await wait(500) // wait for drawer form state to initialize (it's a flake)
-      const uploadListDrawer = page.locator('dialog[id^=list-drawer_1_]').first() // IDs starting with list-drawer_1_ (there's some other symbol after the underscore)
-      await expect(uploadListDrawer).toBeVisible()
-      // find button which has a span with text "payload.jpg" and click it in playwright
-      const uploadButton = uploadListDrawer.locator('button').getByText('payload.jpg').first()
-      await expect(uploadButton).toBeVisible()
-      await uploadButton.click()
-      await expect(uploadListDrawer).toBeHidden()
-      // Check if the upload is there
-      await expect(
-        newSubLexicalAndUploadBlock.locator('.field-type.upload .file-meta__url a'),
-      ).toHaveText('payload.jpg')
+      await expect(async () => {
+        const chooseExistingUploadButton = newSubLexicalAndUploadBlock
+          .locator('.upload__toggler.list-drawer__toggler')
+          .first()
+        await wait(300)
+        await expect(chooseExistingUploadButton).toBeVisible()
+        await wait(300)
+        await chooseExistingUploadButton.click()
+        await wait(500) // wait for drawer form state to initialize (it's a flake)
+        const uploadListDrawer = page.locator('dialog[id^=list-drawer_1_]').first() // IDs starting with list-drawer_1_ (there's some other symbol after the underscore)
+        await expect(uploadListDrawer).toBeVisible()
+        await wait(300)
+
+        // find button which has a span with text "payload.jpg" and click it in playwright
+        const uploadButton = uploadListDrawer.locator('button').getByText('payload.jpg').first()
+        await expect(uploadButton).toBeVisible()
+        await wait(300)
+        await uploadButton.click()
+        await wait(300)
+        await expect(uploadListDrawer).toBeHidden()
+        // Check if the upload is there
+        await expect(
+          newSubLexicalAndUploadBlock.locator('.field-type.upload .file-meta__url a'),
+        ).toHaveText('payload.jpg')
+      }).toPass({
+        timeout: POLL_TOPASS_TIMEOUT,
+      })
+
+      await wait(300)
+
       // save document and assert
       await saveDocAndAssert(page)
+      await wait(300)
+
       await expect(
         newSubLexicalAndUploadBlock.locator('.field-type.upload .file-meta__url a'),
       ).toHaveText('payload.jpg')
       await expect(paragraphInSubEditor).toHaveText('Some subText')
+      await wait(300)
 
       // reload page and assert again
       await page.reload()
+      await wait(300)
+
       await expect(
         newSubLexicalAndUploadBlock.locator('.field-type.upload .file-meta__url a'),
       ).toHaveText('payload.jpg')