fix: rich text editor behaviour in forms

packages/picasso/src/LexicalEditor/LexicalEditor.test.tsx

@@ -6,6 +6,7 @@ import { OnChangePlugin } from '@lexical/react/LexicalOnChangePlugin'
 import LexicalEditor from './LexicalEditor'
 import type { Props } from './LexicalEditor'
 import LexicalTextLengthPlugin from '../LexicalTextLengthPlugin'
+import LexicalHeadingsReplacementPlugin from '../LexicalHeadingsReplacementPlugin'
 import ToolbarPlugin from '../LexicalEditorToolbarPlugin'
 import LexicalListPlugin from '../LexicalListPlugin'
 
@@ -43,9 +44,19 @@ jest.mock('../LexicalTextLengthPlugin', () => ({
   default: jest.fn(() => <div>LexicalTextLengthPlugin</div>),
 }))
 
+jest.mock('../LexicalHeadingsReplacementPlugin', () => ({
+  __esModule: true,
+  default: jest.fn(() => <div>LexicalHeadingsReplacementPlugin</div>),
+}))
+
 const mockedLexicalTextLengthPlugin =
   LexicalTextLengthPlugin as jest.MockedFunction<typeof LexicalTextLengthPlugin>
 
+const mockedLexicalHeadingsReplacementPlugin =
+  LexicalHeadingsReplacementPlugin as jest.MockedFunction<
+    typeof LexicalHeadingsReplacementPlugin
+  >
+
 const mockedToolbarPlugin = ToolbarPlugin as jest.MockedFunction<
   typeof ToolbarPlugin
 >
@@ -67,6 +78,7 @@ const renderLexicalEditor = (props: Partial<OmitInternalProps<Props>> = {}) => {
 describe('LexicalEditor', () => {
   beforeEach(() => {
     mockedLexicalTextLengthPlugin.mockImplementation(() => null)
+    mockedLexicalHeadingsReplacementPlugin.mockImplementation(() => null)
     mockedToolbarPlugin.mockImplementation(() => (
       <div>LexicalEditorToolbarPlugin</div>
     ))
@@ -98,6 +110,14 @@ describe('LexicalEditor', () => {
       )
     })
 
+    it('renders LexicalHeadingsReplacementPlugin', async () => {
+      renderLexicalEditor()
+
+      await waitFor(() =>
+        expect(mockedLexicalHeadingsReplacementPlugin).toHaveBeenCalled()
+      )
+    })
+
     it('renders OnChangePlugin with correct props', () => {
       renderLexicalEditor()
 

packages/picasso/src/LexicalEditor/LexicalEditor.tsx

@@ -24,6 +24,7 @@ import type { ChangeHandler, TextLengthChangeHandler } from './types'
 import ToolbarPlugin from '../LexicalEditorToolbarPlugin'
 import LexicalTextLengthPlugin from '../LexicalTextLengthPlugin'
 import LexicalListPlugin from '../LexicalListPlugin'
+import LexicalHeadingsReplacementPlugin from '../LexicalHeadingsReplacementPlugin'
 
 const useStyles = makeStyles<Theme>(styles, {
   name: 'LexicalEditor',
@@ -173,6 +174,7 @@ const LexicalEditor = forwardRef<HTMLDivElement, Props>(function LexicalEditor(
         <OnChangePlugin ignoreSelectionChange onChange={handleChange} />
         {autoFocus && <AutoFocusPlugin />}
 
+        <LexicalHeadingsReplacementPlugin />
         <LexicalTextLengthPlugin onTextLengthChange={onTextLengthChange} />
         <LexicalListPlugin />
 

packages/picasso/src/LexicalEditor/utils/index.ts

@@ -3,4 +3,3 @@ export { synchronizeToolbarState } from './synchronizeToolbarState'
 export { toolbarStateReducer } from './toolbarState'
 export { getLexicalNode } from './getLexicalNode'
 export { createLexicalTheme } from './createLexicalTheme'
-export { replaceHeadingNodes } from './replaceHeadingNodes'

packages/picasso/src/LexicalEditor/utils/registerLexicalEvents.ts

@@ -1,8 +1,5 @@
-import type { LexicalEditor } from 'lexical'
-import { HeadingNode } from '@lexical/rich-text'
 import { COMMAND_PRIORITY_CRITICAL, SELECTION_CHANGE_COMMAND } from 'lexical'
-
-import { replaceHeadingNodes } from './replaceHeadingNodes'
+import type { LexicalEditor } from 'lexical'
 
 export type LexicalRegisterParams = {
   editor: LexicalEditor
@@ -31,15 +28,9 @@ export const registerLexicalEvents = ({
     COMMAND_PRIORITY_CRITICAL
   )
 
-  const headingListener = editor.registerNodeTransform(
-    HeadingNode,
-    replaceHeadingNodes
-  )
-
   // Cleanup is necessary to avoid listeners piling up with useEffect
   return () => {
     editorListenerCleanup()
     editorCommandsCleanup()
-    headingListener()
   }
 }

packages/picasso/src/LexicalHeadingsReplacementPlugin/LexicalHeadingsReplacementPlugin.ts

@@ -0,0 +1,17 @@
+import { useLexicalComposerContext } from '@lexical/react/LexicalComposerContext'
+import { HeadingNode } from '@lexical/rich-text'
+import { useEffect } from 'react'
+
+import { replaceHeadingNodes } from './utils'
+
+const LexicalHeadingsReplacementPlugin = () => {
+  const [editor] = useLexicalComposerContext()
+
+  useEffect(() => {
+    return editor.registerNodeTransform(HeadingNode, replaceHeadingNodes)
+  }, [editor])
+
+  return null
+}
+
+export default LexicalHeadingsReplacementPlugin

packages/picasso/src/LexicalHeadingsReplacementPlugin/index.ts

@@ -0,0 +1 @@
+export { default } from './LexicalHeadingsReplacementPlugin'

packages/picasso/src/LexicalHeadingsReplacementPlugin/utils/index.ts

@@ -0,0 +1 @@
+export { replaceHeadingNodes } from './replaceHeadingNodes'

packages/picasso/src/LexicalEditor/utils/replaceHeadingNodes.test.ts → packages/picasso/src/LexicalHeadingsReplacementPlugin/utils/replaceHeadingNodes.test.ts

@@ -52,11 +52,13 @@ describe('replaceHeadingNodes', () => {
       )
 
       const appendMock = jest.fn()
+      const selectMock = jest.fn()
 
       mockedCreateParagraphNode.mockImplementation(
         () =>
           ({
             append: appendMock,
+            select: selectMock,
           } as unknown as ParagraphNode)
       )
 
@@ -67,8 +69,12 @@ describe('replaceHeadingNodes', () => {
 
       expect(setFormatMock).toHaveBeenCalledWith('bold')
       expect(appendMock).toHaveBeenCalledWith({ setFormat: setFormatMock })
+      expect(selectMock).toHaveBeenCalled()
 
-      expect(node.replace).toHaveBeenCalledWith({ append: appendMock })
+      expect(node.replace).toHaveBeenCalledWith({
+        append: appendMock,
+        select: selectMock,
+      })
     })
   })
 })