fix(ai): edit previous message parse in rich editor

- Integrated ScrollArea for better message editing experience in the EditableMessage component.
- Updated initial editor state handling to streamline the setup process.
- Added MentionPlugin to support mentions within the editor, enhancing user interaction.
- Refactored LexicalRichEditor to utilize EditorRefPlugin for improved editor reference management.

These changes improve the usability and functionality of the message editing interface in the AI chat application.

Signed-off-by: Innei <tukon479@gmail.com>

apps/desktop/layer/renderer/src/modules/ai-chat/components/message/EditableMessage.tsx

@@ -1,9 +1,7 @@
 import type { LexicalRichEditorRef } from "@follow/components/ui/lexical-rich-editor/index.js"
-import {
-  createDefaultLexicalEditor,
-  LexicalRichEditor,
-} from "@follow/components/ui/lexical-rich-editor/index.js"
-import { cn } from "@follow/utils"
+import { LexicalRichEditor } from "@follow/components/ui/lexical-rich-editor/index.js"
+import { ScrollArea } from "@follow/components/ui/scroll-area/ScrollArea.js"
+import { cn, nextFrame } from "@follow/utils"
 import type { BizUIMessage } from "@folo-services/ai-tools"
 import { isEqual } from "es-toolkit"
 import type { EditorState, LexicalEditor, SerializedEditorState } from "lexical"
@@ -13,6 +11,8 @@ import { useCallback, useEffect, useMemo, useRef, useState } from "react"
 import { useEditingMessageId, useSetEditingMessageId } from "~/modules/ai-chat/atoms/session"
 import { useChatStatus } from "~/modules/ai-chat/store/hooks"
 
+import { MentionPlugin } from "../../editor"
+
 interface EditableMessageProps {
   messageId: string
   parts: BizUIMessage["parts"]
@@ -36,10 +36,10 @@ export const EditableMessage = ({
   const [currentEditor, setCurrentEditor] = useState<LexicalEditor | null>(null)
 
   const initialEditorState = useMemo(() => {
-    const serializedEditorState = (parts.find((part) => part.type === "data-rich-text") as any)
-      ?.data.state as SerializedEditorState
-    return createDefaultLexicalEditor().parseEditorState(serializedEditorState)
-  }, [parts])
+    return (parts.find((part) => part.type === "data-rich-text") as any)?.data
+      .state as SerializedEditorState
+  }, [])
+
   const isEditing = editingMessageId === messageId
   const isProcessing = status === "submitted" || status === "streaming"
 
@@ -49,13 +49,26 @@ export const EditableMessage = ({
       // Focus the editor
       editorRef.current.focus()
     }
-  }, [isEditing, initialEditorState, currentEditor])
+  }, [isEditing, currentEditor])
+
+  const setInitialEditorStateOnceRef = useRef(false)
+
+  useEffect(() => {
+    return nextFrame(() => {
+      if (setInitialEditorStateOnceRef.current) return
+      const editor = editorRef.current?.getEditor()
+
+      if (!editor) return
+      editor.setEditorState(editor.parseEditorState(initialEditorState))
+      setInitialEditorStateOnceRef.current = true
+    })
+  }, [initialEditorState])
 
   const handleSave = useCallback(() => {
     if (currentEditor && editorRef.current && !editorRef.current.isEmpty()) {
       const serializedEditorState = currentEditor.getEditorState().toJSON()
 
-      if (!isEqual(serializedEditorState, initialEditorState.toJSON())) {
+      if (!isEqual(serializedEditorState, initialEditorState)) {
         onSave(serializedEditorState, currentEditor)
       }
     }
@@ -102,15 +115,17 @@ export const EditableMessage = ({
     <div className={cn("relative", className)}>
       {/* Edit input */}
       <div className="bg-background/60 focus-within:ring-accent/20 focus-within:border-accent/80 border-border/80 relative overflow-hidden rounded-xl border backdrop-blur-xl duration-200 focus-within:ring-2">
-        <LexicalRichEditor
-          ref={editorRef}
-          placeholder="Edit your message..."
-          initalEditorState={initialEditorState}
-          className="w-full pr-20"
-          onChange={handleEditorChange}
-          onKeyDown={handleKeyDown}
-          namespace="EditableMessageRichEditor"
-        />
+        <ScrollArea rootClassName="mx-5 my-3.5 mr-20 flex-1 overflow-auto">
+          <LexicalRichEditor
+            ref={editorRef}
+            placeholder="Edit your message..."
+            className="w-full min-w-64"
+            onChange={handleEditorChange}
+            onKeyDown={handleKeyDown}
+            namespace="EditableMessageRichEditor"
+            plugins={[MentionPlugin]}
+          />
+        </ScrollArea>
 
         {/* Action buttons */}
         <div className="absolute right-2 top-1/2 flex -translate-y-1/2 gap-1">

packages/internal/components/src/ui/lexical-rich-editor/LexicalRichEditor.tsx

@@ -4,6 +4,7 @@ import { AutoFocusPlugin } from "@lexical/react/LexicalAutoFocusPlugin"
 import type { InitialConfigType } from "@lexical/react/LexicalComposer"
 import { LexicalComposer } from "@lexical/react/LexicalComposer"
 import { ContentEditable } from "@lexical/react/LexicalContentEditable"
+import { EditorRefPlugin } from "@lexical/react/LexicalEditorRefPlugin"
 import { LexicalErrorBoundary } from "@lexical/react/LexicalErrorBoundary"
 import { HistoryPlugin } from "@lexical/react/LexicalHistoryPlugin"
 import { LinkPlugin } from "@lexical/react/LexicalLinkPlugin"
@@ -13,7 +14,7 @@ import { OnChangePlugin } from "@lexical/react/LexicalOnChangePlugin"
 import { RichTextPlugin } from "@lexical/react/LexicalRichTextPlugin"
 import type { EditorState, LexicalEditor } from "lexical"
 import { $getRoot } from "lexical"
-import { useImperativeHandle, useRef, useState } from "react"
+import { useCallback, useImperativeHandle, useState } from "react"
 
 import { LexicalRichEditorNodes } from "./nodes"
 import { KeyboardPlugin } from "./plugins"
@@ -44,7 +45,7 @@ export const LexicalRichEditor = ({
   initalEditorState,
   plugins,
 }: LexicalRichEditorProps & { ref?: React.RefObject<LexicalRichEditorRef | null> }) => {
-  const editorRef = useRef<LexicalEditor | null>(null)
+  const [editorRef, setEditorRef] = useState<LexicalEditor | null>(null)
   const [isEmpty, setIsEmpty] = useState(true)
 
   // Collect nodes from plugins
@@ -61,23 +62,27 @@ export const LexicalRichEditor = ({
     editorState: initalEditorState,
   }
 
-  useImperativeHandle(ref, () => ({
-    getEditor: () => editorRef.current!,
-    focus: () => {
-      editorRef.current?.focus()
-    },
-    clear: () => {
-      editorRef.current?.update(() => {
-        const root = $getRoot()
-        root.clear()
-      })
-    },
-    isEmpty: () => isEmpty,
-  }))
+  useImperativeHandle(
+    ref,
+    useCallback(
+      () => ({
+        getEditor: () => editorRef!,
+        focus: () => {
+          editorRef?.focus()
+        },
+        clear: () => {
+          editorRef?.update(() => {
+            const root = $getRoot()
+            root.clear()
+          })
+        },
+        isEmpty: () => isEmpty,
+      }),
+      [isEmpty, editorRef],
+    ),
+  )
 
   const handleChange = (editorState: EditorState, editor: LexicalEditor) => {
-    editorRef.current = editor
-
     // Check if editor is empty
     editorState.read(() => {
       const root = $getRoot()
@@ -96,7 +101,7 @@ export const LexicalRichEditor = ({
             <ContentEditable
               className={cn(
                 "scrollbar-none text-text placeholder:text-text-secondary cursor-text",
-                "max-h-40 min-h-14 w-full resize-none bg-transparent",
+                "h-14 w-full resize-none bg-transparent",
                 "text-sm !outline-none transition-all duration-200 focus:outline-none",
               )}
               aria-placeholder={placeholder}
@@ -110,6 +115,7 @@ export const LexicalRichEditor = ({
           ErrorBoundary={LexicalErrorBoundary}
         />
         <OnChangePlugin onChange={handleChange} />
+        <EditorRefPlugin editorRef={setEditorRef} />
 
         {enabledPlugins.history && <HistoryPlugin />}
         {enabledPlugins.markdown && <MarkdownShortcutPlugin transformers={TRANSFORMERS} />}

packages/internal/utils/src/dom.ts

@@ -4,11 +4,16 @@ export const stopPropagation = <T extends { stopPropagation: () => any }>(e: T)
 export const preventDefault = <T extends { preventDefault: () => any }>(e: T) => e.preventDefault()
 
 export const nextFrame = (fn: (...args: any[]) => any) => {
-  requestAnimationFrame(() => {
-    requestAnimationFrame(() => {
+  let timer2: number | null = null
+  const timer1 = requestAnimationFrame(() => {
+    timer2 = requestAnimationFrame(() => {
       fn()
     })
   })
+  return () => {
+    if (timer1) cancelAnimationFrame(timer1)
+    if (timer2) cancelAnimationFrame(timer2)
+  }
 }
 
 export const getElementTop = (element: HTMLElement) => {