refactor: optimize chat components and remove unused GridContainer

- Refactored ToolInvocationComponent to use React.memo for performance optimization. - Updated AIDisplayAnalyticsPart, AIDisplayEntriesPart, AIDisplayFeedsPart, and AIDisplaySubscriptionsPart to utilize memoization. - Replaced GridContainer with standard div elements for layout consistency. - Removed unused GridContainer and StatisticsOverview components to streamline the codebase. Signed-off-by: Innei <tukon479@gmail.com>
2025-07-31 23:12:51 +08:00 · 2025-07-31 23:12:51 +08:00 · c5ff53aa20
parent 73a1e6eb6c
commit c5ff53aa20
18 changed files with 301 additions and 356 deletions
--- a/apps/desktop/layer/renderer/src/modules/ai/chat/components/ToolInvocationComponent.tsx
+++ b/apps/desktop/layer/renderer/src/modules/ai/chat/components/ToolInvocationComponent.tsx
@ -12,51 +12,54 @@ interface ToolInvocationComponentProps {
  part: ToolUIPart
 }

-export const ToolInvocationComponent: React.FC<ToolInvocationComponentProps> = ({ part }) => {
-  const toolName = getToolName(part)
-  return (
-    <div className="bg-material-medium border-border size-full min-w-0 max-w-prose rounded-lg border text-left">
-      <div className="w-[9999px] max-w-[calc(var(--ai-chat-layout-width,65ch)_-120px)]" />
-      <Accordion type="single" collapsible>
-        <AccordionItem value="tool-invocation">
-          <AccordionTrigger className="flex w-full cursor-pointer items-center gap-3 py-1 pl-4 pr-2 hover:no-underline">
-            {/* Tool Info */}
-            <div className="flex h-6 min-w-0 flex-1 items-center">
-              <div className="flex items-center gap-2 text-xs">
-                <i className="i-mingcute-tool-line" />
-                <span className="text-text-secondary">Tool Calling:</span>
-                <h4 className="text-text truncate font-medium">{toolName}</h4>
+export const ToolInvocationComponent: React.FC<ToolInvocationComponentProps> = React.memo(
+  ({ part }) => {
+    const toolName = getToolName(part)
+
+    return (
+      <div className="bg-material-medium border-border size-full min-w-0 max-w-prose rounded-lg border text-left">
+        <div className="w-[9999px] max-w-[calc(var(--ai-chat-layout-width,65ch)_-120px)]" />
+        <Accordion type="single" collapsible>
+          <AccordionItem value="tool-invocation">
+            <AccordionTrigger className="flex w-full cursor-pointer items-center gap-3 py-1 pl-4 pr-2 hover:no-underline">
+              {/* Tool Info */}
+              <div className="flex h-6 min-w-0 flex-1 items-center">
+                <div className="flex items-center gap-2 text-xs">
+                  <i className="i-mingcute-tool-line" />
+                  <span className="text-text-secondary">Tool Calling:</span>
+                  <h4 className="text-text truncate font-medium">{toolName}</h4>
+                </div>
              </div>
-            </div>
-          </AccordionTrigger>
+            </AccordionTrigger>

-          <AccordionContent className="min-w-0 border-t border-zinc-200/50 bg-zinc-50/50 p-4 dark:border-zinc-700/50 dark:bg-zinc-800/50">
-            <div className="space-y-3">
-              {"input" in part && (
-                <div>
-                  <div className="text-text-tertiary mb-2 text-xs font-semibold uppercase tracking-wide">
-                    Arguments
+            <AccordionContent className="min-w-0 border-t border-zinc-200/50 bg-zinc-50/50 p-4 dark:border-zinc-700/50 dark:bg-zinc-800/50">
+              <div className="space-y-3">
+                {"input" in part && (
+                  <div>
+                    <div className="text-text-tertiary mb-2 text-xs font-semibold uppercase tracking-wide">
+                      Arguments
+                    </div>
+                    <pre className="text-text-secondary max-h-32 overflow-auto rounded-lg bg-zinc-100/80 p-3 text-xs leading-relaxed dark:bg-zinc-900/80">
+                      {JSON.stringify(part.input, null, 2)}
+                    </pre>
                  </div>
-                  <pre className="text-text-secondary max-h-32 overflow-auto rounded-lg bg-zinc-100/80 p-3 text-xs leading-relaxed dark:bg-zinc-900/80">
-                    {JSON.stringify(part.input, null, 2)}
-                  </pre>
-                </div>
-              )}
+                )}

-              {"output" in part && (
-                <div>
-                  <div className="text-text-tertiary mb-2 text-xs font-semibold uppercase tracking-wide">
-                    Result
+                {"output" in part && (
+                  <div>
+                    <div className="text-text-tertiary mb-2 text-xs font-semibold uppercase tracking-wide">
+                      Result
+                    </div>
+                    <pre className="text-text max-h-32 overflow-auto rounded-lg bg-zinc-100/80 p-3 text-xs leading-relaxed dark:bg-zinc-900/80">
+                      {JSON.stringify(part.output, null, 2)}
+                    </pre>
                  </div>
-                  <pre className="text-text max-h-32 overflow-auto rounded-lg bg-zinc-100/80 p-3 text-xs leading-relaxed dark:bg-zinc-900/80">
-                    {JSON.stringify(part.output, null, 2)}
-                  </pre>
-                </div>
-              )}
-            </div>
-          </AccordionContent>
-        </AccordionItem>
-      </Accordion>
-    </div>
-  )
-}
+                )}
+              </div>
+            </AccordionContent>
+          </AccordionItem>
+        </Accordion>
+      </div>
+    )
+  },
+)
--- a/apps/desktop/layer/renderer/src/modules/ai/chat/components/displays/AIDisplayAnalyticsPart.tsx
+++ b/apps/desktop/layer/renderer/src/modules/ai/chat/components/displays/AIDisplayAnalyticsPart.tsx
@ -14,6 +14,7 @@ import {
  TableRow,
 } from "@follow/components/ui/table/index.js"
 import dayjs from "dayjs"
+import { memo } from "react"

 import type { AIDisplayAnalyticsTool } from "../../__internal__/types"
 import { ErrorState, LoadingState } from "../common-states"
@ -229,18 +230,7 @@ const OverviewAnalytics = ({ data }: { data: AnalyticsData["overviewStats"] }) =
  )
 }

-export const AIDisplayAnalyticsPart = ({ part }: { part: AIDisplayAnalyticsTool }) => {
-  // Handle loading state
-  if (part.state === "input-streaming" || part.state === "input-available") {
-    return (
-      <LoadingState
-        title="Loading Analytics..."
-        description="Fetching analytics data..."
-        maxWidth="max-w-4xl"
-      />
-    )
-  }
-
+export const AIDisplayAnalyticsPart = memo(({ part }: { part: AIDisplayAnalyticsTool }) => {
  // Handle error state
  if (part.state === "output-error") {
    return (
@ -256,9 +246,8 @@ export const AIDisplayAnalyticsPart = ({ part }: { part: AIDisplayAnalyticsTool
    )
  }

-  // Handle no output
-  const { output } = part
-  if (!output) {
+  // Handle no output or invalid state
+  if (part.state !== "output-available" || !part.output) {
    return (
      <LoadingState
        title="Loading Analytics..."
@ -268,7 +257,7 @@ export const AIDisplayAnalyticsPart = ({ part }: { part: AIDisplayAnalyticsTool
    )
  }

-  const { analyticsData, analyticsType, timeRange, displayType, title } = output
+  const { analyticsData, analyticsType, timeRange, displayType, title } = part.output

  const renderAnalytics = () => {
    switch (analyticsType) {
@ -309,4 +298,4 @@ export const AIDisplayAnalyticsPart = ({ part }: { part: AIDisplayAnalyticsTool
      <CardContent>{renderAnalytics()}</CardContent>
    </Card>
  )
-}
+})
--- a/apps/desktop/layer/renderer/src/modules/ai/chat/components/displays/AIDisplayEntriesPart.tsx
+++ b/apps/desktop/layer/renderer/src/modules/ai/chat/components/displays/AIDisplayEntriesPart.tsx
@ -14,7 +14,7 @@ import { FeedIcon } from "~/modules/feed/feed-icon"

 import type { AIDisplayEntriesTool } from "../../__internal__/types"
 import { ErrorState, LoadingState } from "../common-states"
-import { CategoryTag, EmptyState, GridContainer, StatCard } from "./shared"
+import { CategoryTag, EmptyState, StatCard } from "./shared"

 type EntryData = AIDisplayEntriesTool["output"]["entries"]
 type EntryItem = EntryData[number]
@ -59,10 +59,7 @@ const EntriesGrid = ({
  }

  return (
-    <GridContainer
-      columns={{ base: 1, md: 2, lg: 3 }}
-      className="@[600px]:grid-cols-2 @[900px]:grid-cols-3"
-    >
+    <div className="@[600px]:grid-cols-2 @[900px]:grid-cols-3 grid grid-cols-1 gap-4 md:grid-cols-2 lg:grid-cols-3">
      {data.map((item) => (
        <EntryCard
          key={item.entry.id}
@ -71,7 +68,7 @@ const EntriesGrid = ({
          showMetadata={showMetadata}
        />
      ))}
-    </GridContainer>
+    </div>
  )
 }

@ -266,18 +263,7 @@ const GroupedEntries = ({
  )
 }

-export const AIDisplayEntriesPart = ({ part }: { part: AIDisplayEntriesTool }) => {
-  // Handle loading state
-  if (part.state === "input-streaming" || part.state === "input-available") {
-    return (
-      <LoadingState
-        title="Loading Entries..."
-        description="Fetching entry data..."
-        maxWidth="max-w-6xl"
-      />
-    )
-  }
-
+export const AIDisplayEntriesPart = memo(({ part }: { part: AIDisplayEntriesTool }) => {
  // Handle error state
  if (part.state === "output-error") {
    return (
@ -355,16 +341,16 @@ export const AIDisplayEntriesPart = ({ part }: { part: AIDisplayEntriesTool }) =
      </CardHeader>
      <CardContent className="@container space-y-6">
        {/* Statistics Overview */}
-        <GridContainer columns={{ base: 2, md: 4 }} className="@[600px]:grid-cols-4">
+        <div className="@[600px]:grid-cols-4 grid grid-cols-2 gap-4 md:grid-cols-4">
          <StatCard title="Total Entries" value={totalEntries} emoji="📄" />
          <StatCard title="Feeds" value={feedsCount} emoji="📡" />
          <StatCard title="Authors" value={authorsCount} emoji="✍️" />
          <StatCard title="Categories" value={categoriesCount} emoji="🏷️" />
-        </GridContainer>
+        </div>

        {/* Entries Display */}
        {renderEntries()}
      </CardContent>
    </Card>
  )
-}
+})
--- a/apps/desktop/layer/renderer/src/modules/ai/chat/components/displays/AIDisplayFeedsPart.tsx
+++ b/apps/desktop/layer/renderer/src/modules/ai/chat/components/displays/AIDisplayFeedsPart.tsx
@ -6,12 +6,13 @@ import {
  CardTitle,
 } from "@follow/components/ui/card/index.js"
 import dayjs from "dayjs"
+import { memo } from "react"

 import { FeedIcon } from "~/modules/feed/feed-icon"

 import type { AIDisplayFeedsTool } from "../../__internal__/types"
 import { ErrorState, LoadingState } from "../common-states"
-import { AnalyticsMetrics, EmptyState, GridContainer, StatCard } from "./shared"
+import { AnalyticsMetrics, EmptyState, StatCard } from "./shared"

 type FeedData = AIDisplayFeedsTool["output"]["feeds"]

@ -21,7 +22,7 @@ const FeedsGrid = ({ data, showAnalytics }: { data: FeedData; showAnalytics: boo
  }

  return (
-    <GridContainer columns={{ base: 2, md: 3 }} className="@[600px]:grid-cols-3">
+    <div className="@[600px]:grid-cols-3 grid grid-cols-2 gap-4 md:grid-cols-3">
      {data.map((item) => (
        <Card key={item.feed.id} className="p-4">
          <CardHeader className="h-24 px-2 py-3">
@ -72,22 +73,11 @@ const FeedsGrid = ({ data, showAnalytics }: { data: FeedData; showAnalytics: boo
          </CardContent>
        </Card>
      ))}
-    </GridContainer>
+    </div>
  )
 }

-export const AIDisplayFeedsPart = ({ part }: { part: AIDisplayFeedsTool }) => {
-  // Handle loading state
-  if (part.state === "input-streaming" || part.state === "input-available") {
-    return (
-      <LoadingState
-        title="Loading Feeds..."
-        description="Fetching feed data..."
-        maxWidth="max-w-6xl"
-      />
-    )
-  }
-
+export const AIDisplayFeedsPart = memo(({ part }: { part: AIDisplayFeedsTool }) => {
  // Handle error state
  if (part.state === "output-error") {
    return (
@ -145,7 +135,7 @@ export const AIDisplayFeedsPart = ({ part }: { part: AIDisplayFeedsTool }) => {
      </CardHeader>
      <CardContent className="@container space-y-6">
        {/* Statistics Overview */}
-        <GridContainer columns={{ base: 2, md: 4 }} className="@[600px]:grid-cols-4">
+        <div className="@[600px]:grid-cols-4 grid grid-cols-2 gap-4 md:grid-cols-4">
          <StatCard title="Total Feeds" value={totalFeeds} emoji="📊" />
          <StatCard
            title="Active Feeds"
@ -163,11 +153,11 @@ export const AIDisplayFeedsPart = ({ part }: { part: AIDisplayFeedsTool }) => {
              <StatCard title="Total Views" value={totalViews.toLocaleString()} emoji="👀" />
            </>
          )}
-        </GridContainer>
+        </div>

        {/* Feeds Display */}
        {renderFeeds()}
      </CardContent>
    </Card>
  )
-}
+})
--- a/apps/desktop/layer/renderer/src/modules/ai/chat/components/displays/AIDisplaySubscriptionsPart.tsx
+++ b/apps/desktop/layer/renderer/src/modules/ai/chat/components/displays/AIDisplaySubscriptionsPart.tsx
@ -6,12 +6,13 @@ import {
  CardTitle,
 } from "@follow/components/ui/card/index.js"
 import dayjs from "dayjs"
+import { memo } from "react"

 import { FeedIcon } from "~/modules/feed/feed-icon"

 import type { AIDisplaySubscriptionsTool } from "../../__internal__/types"
 import { ErrorState, LoadingState } from "../common-states"
-import { AnalyticsMetrics, CategoryTag, EmptyState, GridContainer, StatCard } from "./shared"
+import { AnalyticsMetrics, CategoryTag, EmptyState, StatCard } from "./shared"

 type SubscriptionData = AIDisplaySubscriptionsTool["output"]["subscriptions"]

@ -58,7 +59,7 @@ const SubscriptionsGrid = ({
  }

  return (
-    <GridContainer columns={{ base: 2, md: 3 }} className="@[600px]:grid-cols-3">
+    <div className="@[600px]:grid-cols-3 @[400px]:grid-cols-2 grid grid-cols-1 gap-4">
      {data.map((sub) => (
        <Card key={`${sub.subscription?.userId}-${sub.subscription?.feedId}`} className="p-4">
          <CardHeader className="h-24 px-2 py-3">
@ -105,7 +106,7 @@ const SubscriptionsGrid = ({
          </CardContent>
        </Card>
      ))}
-    </GridContainer>
+    </div>
  )
 }

@ -171,18 +172,8 @@ const GroupedSubscriptions = ({
  )
 }

-export const AIDisplaySubscriptionsPart = ({ part }: { part: AIDisplaySubscriptionsTool }) => {
-  // Handle loading state
-  if (part.state === "input-streaming" || part.state === "input-available") {
-    return (
-      <LoadingState
-        title="Loading Subscriptions..."
-        description="Fetching subscription data..."
-        maxWidth="max-w-6xl"
-      />
-    )
-  }
-
+export const AIDisplaySubscriptionsPart = memo(({ part }: { part: AIDisplaySubscriptionsTool }) => {
+  return null
  // Handle error state
  if (part.state === "output-error") {
    return (
@ -194,7 +185,7 @@ export const AIDisplaySubscriptionsPart = ({ part }: { part: AIDisplaySubscripti
    )
  }

-  // Handle no output or invalid state
+  // Handle loading state
  if (part.state !== "output-available" || !part.output) {
    return (
      <LoadingState
@ -266,7 +257,7 @@ export const AIDisplaySubscriptionsPart = ({ part }: { part: AIDisplaySubscripti
      </CardHeader>
      <CardContent className="@container space-y-6">
        {/* Statistics Overview */}
-        <GridContainer columns={{ base: 2, md: 4 }} className="@[600px]:grid-cols-4">
+        <div className="@[700px]:grid-cols-4 grid grid-cols-2 gap-4">
          <StatCard title="Total Subscriptions" value={totalSubscriptions} emoji="📊" />
          <StatCard
            title="Active Feeds"
@ -276,11 +267,11 @@ export const AIDisplaySubscriptionsPart = ({ part }: { part: AIDisplaySubscripti
          />
          {showCategories && <StatCard title="Categories" value={categoriesCount} emoji="🏷️" />}
          <StatCard title="Total Views" value={totalViews.toLocaleString()} emoji="👀" />
-        </GridContainer>
+        </div>

        {/* Subscriptions Display */}
        {renderSubscriptions()}
      </CardContent>
    </Card>
  )
-}
+})
--- a/apps/desktop/layer/renderer/src/modules/ai/chat/components/displays/shared/GridContainer.tsx
+++ b/apps/desktop/layer/renderer/src/modules/ai/chat/components/displays/shared/GridContainer.tsx
@ -1,33 +0,0 @@
-import { cn } from "@follow/utils/utils"
-import type { ReactNode } from "react"
-
-export interface GridContainerProps {
-  columns?: {
-    base: number
-    md: number
-    lg?: number
-  }
-  gap?: number
-  children: ReactNode
-  className?: string
-}
-
-export const GridContainer = ({
-  columns = { base: 1, md: 2 },
-  gap = 4,
-  children,
-  className,
-}: GridContainerProps) => (
-  <div
-    className={cn(
-      "grid",
-      `grid-cols-${columns.base}`,
-      `md:grid-cols-${columns.md}`,
-      columns.lg && `lg:grid-cols-${columns.lg}`,
-      `gap-${gap}`,
-      className,
-    )}
-  >
-    {children}
-  </div>
-)
--- a/apps/desktop/layer/renderer/src/modules/ai/chat/components/displays/shared/StatisticsOverview.tsx
+++ b/apps/desktop/layer/renderer/src/modules/ai/chat/components/displays/shared/StatisticsOverview.tsx
@ -1,34 +0,0 @@
-import { cn } from "@follow/utils/utils"
-
-import type { StatCardProps } from "./StatCard"
-import { StatCard } from "./StatCard"
-
-export interface StatisticsOverviewProps {
-  stats: StatCardProps[]
-  columns?: {
-    base: number
-    md: number
-    lg?: number
-  }
-  className?: string
-}
-
-export const StatisticsOverview = ({
-  stats,
-  columns = { base: 2, md: 4 },
-  className,
-}: StatisticsOverviewProps) => (
-  <div
-    className={cn(
-      "grid gap-4",
-      `grid-cols-${columns.base}`,
-      `md:grid-cols-${columns.md}`,
-      columns.lg && `lg:grid-cols-${columns.lg}`,
-      className,
-    )}
-  >
-    {stats.map((stat, index) => (
-      <StatCard key={index} {...stat} />
-    ))}
-  </div>
-)
--- a/apps/desktop/layer/renderer/src/modules/ai/chat/components/displays/shared/index.ts
+++ b/apps/desktop/layer/renderer/src/modules/ai/chat/components/displays/shared/index.ts
@ -4,7 +4,5 @@ export * from "./ChartPlaceholder"
 export * from "./DisplayHeader"
 export * from "./EmptyState"
 export * from "./FeedItemCard"
-export * from "./GridContainer"
 export * from "./GroupedContent"
 export * from "./StatCard"
-export * from "./StatisticsOverview"
--- a/apps/desktop/layer/renderer/src/modules/ai/chat/components/layouts/AIChatSendButton.tsx
+++ b/apps/desktop/layer/renderer/src/modules/ai/chat/components/layouts/AIChatSendButton.tsx
@ -26,7 +26,7 @@ export const AIChatSendButton: FC<AIChatSendButtonProps> = ({
          ? "bg-red-500/90 hover:bg-red-500 shadow-lg shadow-red-500/25 backdrop-blur-sm"
          : disabled
            ? "bg-gray-200/80 cursor-not-allowed backdrop-blur-sm"
-            : "bg-gradient-to-r from-blue-500/90 to-blue-600/90 hover:from-blue-500 hover:to-blue-600 shadow-lg shadow-blue-500/25 backdrop-blur-sm hover:shadow-blue-500/35",
+            : "bg-gradient-to-r from-accent to-accent/90 hover:from-accent hover:to-accent/90 shadow-lg shadow-accent/25 backdrop-blur-sm hover:shadow-accent/35",
        className,
      )}
    >
--- a/apps/desktop/layer/renderer/src/modules/ai/chat/components/message/AIChatMessage.tsx
+++ b/apps/desktop/layer/renderer/src/modules/ai/chat/components/message/AIChatMessage.tsx
@ -1,5 +1,7 @@
+import { createDefaultLexicalEditor } from "@follow/components/ui/lexical-rich-editor/editor.js"
 import { stopPropagation } from "@follow/utils"
 import type { UIDataTypes, UIMessage } from "ai"
+import type { LexicalEditor, SerializedEditorState } from "lexical"
 import { m } from "motion/react"
 import * as React from "react"
 import { toast } from "sonner"
@ -9,6 +11,8 @@ import { useChatActions } from "~/modules/ai/chat/__internal__/hooks"
 import type { BizUIMetadata, BizUITools } from "~/modules/ai/chat/__internal__/types"
 import { useEditingMessageId, useSetEditingMessageId } from "~/modules/ai/chat/atoms/session"

+import type { RichTextPart } from "../../types/ChatSession"
+import { convertLexicalToMarkdown } from "../../utils/lexical-markdown"
 import { AIMessageParts } from "./AIMessageParts"
 import { EditableMessage } from "./EditableMessage"

@ -36,39 +40,67 @@ export const AIChatMessage: React.FC<AIChatMessageProps> = React.memo(({ message
  const isEditing = editingMessageId === messageId
  const isUserMessage = message.role === "user"

-  // Get message content for editing
-  const messageContent = React.useMemo(() => {
-    return (
-      message.parts
-        ?.filter((part) => part.type === "text")
-        .map((part) => part.text)
-        .join(" ") || ""
-    )
-  }, [message.parts])
-
  const handleEdit = React.useCallback(() => {
    if (isUserMessage) {
      setEditingMessageId(messageId)
    }
  }, [isUserMessage, messageId, setEditingMessageId])

-  const handleSaveEdit = React.useCallback(
-    (newContent: string) => {
-      const messages = chatActions.getMessages()
-      if (newContent.trim() !== messageContent.trim()) {
-        // Find the message index and remove all messages after it (including AI responses)
-        const messageIndex = messages.findIndex((msg) => msg.id === messageId)
-        if (messageIndex !== -1) {
-          const messagesToKeep = messages.slice(0, messageIndex)
-          chatActions.setMessages(messagesToKeep)
-
-          // Send the edited message
-          chatActions.sendMessage(newContent)
+  const getMessageMarkdownFormat = React.useCallback(() => {
+    let content = ""
+    for (const part of message.parts) {
+      let lexicalEditor: LexicalEditor | null = null
+      switch (part.type) {
+        case "text": {
+          content += part.text
+          break
        }
+        case "data-rich-text": {
+          lexicalEditor ||= createDefaultLexicalEditor()
+          lexicalEditor.setEditorState(
+            lexicalEditor.parseEditorState((part as RichTextPart).data.state),
+          )
+          content += convertLexicalToMarkdown(lexicalEditor)
+          break
+        }
+
+        default: {
+          if (part.type.startsWith("tool-")) {
+            content += `\n\n[TOOL CALL: ${part.type.replace("tool-", "")}]\n\n`
+          }
+          break
+        }
+      }
+    }
+    return content
+  }, [message.parts])
+
+  const handleSaveEdit = React.useCallback(
+    (newState: SerializedEditorState, editor: LexicalEditor) => {
+      const messageContent = convertLexicalToMarkdown(editor)
+      const messages = chatActions.getMessages()
+      const messageIndex = messages.findIndex((msg) => msg.id === messageId)
+      if (messageIndex !== -1) {
+        const messagesToKeep = messages.slice(0, messageIndex)
+        const nextMessage = messages[messageIndex]!
+        chatActions.setMessages(messagesToKeep)
+
+        const richTextPart = nextMessage.parts.find(
+          (part) => part.type === "data-rich-text",
+        ) as RichTextPart
+        if (richTextPart) {
+          richTextPart.data = {
+            state: newState,
+            text: messageContent,
+          }
+        }
+
+        // Send the edited message
+        chatActions.sendMessage(nextMessage)
      }
      setEditingMessageId(null)
    },
-    [messageContent, messageId, chatActions, setEditingMessageId],
+    [chatActions, messageId, setEditingMessageId],
  )

  const handleCancelEdit = React.useCallback(() => {
@ -76,13 +108,14 @@ export const AIChatMessage: React.FC<AIChatMessageProps> = React.memo(({ message
  }, [setEditingMessageId])

  const handleCopy = React.useCallback(async () => {
+    const messageContent = getMessageMarkdownFormat()
    try {
      await copyToClipboard(messageContent)
      toast.success("Message copied to clipboard")
    } catch {
      toast.error("Failed to copy message")
    }
-  }, [messageContent])
+  }, [getMessageMarkdownFormat])

  const handleRetry = React.useCallback(() => {
    chatActions.regenerate({ messageId })
@ -115,7 +148,7 @@ export const AIChatMessage: React.FC<AIChatMessageProps> = React.memo(({ message
        {isEditing && isUserMessage ? (
          <EditableMessage
            messageId={messageId}
-            initialContent={messageContent}
+            parts={message.parts}
            onSave={handleSaveEdit}
            onCancel={handleCancelEdit}
          />
--- a/apps/desktop/layer/renderer/src/modules/ai/chat/components/message/AIRichTextMessage.tsx
+++ b/apps/desktop/layer/renderer/src/modules/ai/chat/components/message/AIRichTextMessage.tsx
@ -1,17 +1,12 @@
-import { defaultLexicalTheme } from "@follow/components/ui/lexical-rich-editor/index.js"
+import { LexicalRichEditorNodes } from "@follow/components/ui/lexical-rich-editor/nodes.js"
+import { defaultLexicalTheme } from "@follow/components/ui/lexical-rich-editor/theme.js"
 import { cn } from "@follow/utils"
-import { CodeHighlightNode, CodeNode } from "@lexical/code"
-import { LinkNode } from "@lexical/link"
-import { ListItemNode, ListNode } from "@lexical/list"
-import { MarkNode } from "@lexical/mark"
 import type { InitialConfigType } from "@lexical/react/LexicalComposer"
 import { LexicalComposer } from "@lexical/react/LexicalComposer"
 import { ContentEditable } from "@lexical/react/LexicalContentEditable"
 import { LexicalErrorBoundary } from "@lexical/react/LexicalErrorBoundary"
 import { RichTextPlugin } from "@lexical/react/LexicalRichTextPlugin"
-import { HeadingNode, QuoteNode } from "@lexical/rich-text"
 import type { SerializedEditorState } from "lexical"
-import { ParagraphNode, TextNode } from "lexical"
 import * as React from "react"

 function onError(error: Error) {
@ -26,49 +21,29 @@ interface AIRichTextMessageProps {
  className?: string
 }

-export const AIRichTextMessage: React.FC<AIRichTextMessageProps> = ({ data, className }) => {
-  const initialConfig: InitialConfigType = {
-    namespace: "AIRichTextDisplay",
-    theme: defaultLexicalTheme,
-    onError,
-    editable: false, // Read-only mode
-    editorState: JSON.stringify(data.state),
-    nodes: [
-      // Core nodes
-      ParagraphNode,
-      TextNode,
+export const AIRichTextMessage: React.FC<AIRichTextMessageProps> = React.memo(
+  ({ data, className }) => {
+    const initialConfig: InitialConfigType = {
+      namespace: "AIRichTextDisplay",
+      theme: defaultLexicalTheme,
+      onError,
+      editable: false, // Read-only mode
+      editorState: JSON.stringify(data.state),
+      nodes: LexicalRichEditorNodes,
+    }

-      // Rich text nodes
-      HeadingNode,
-      QuoteNode,
-
-      // List nodes
-      ListNode,
-      ListItemNode,
-
-      // Code nodes
-      CodeNode,
-      CodeHighlightNode,
-
-      // Link nodes
-      LinkNode,
-
-      // Text format nodes
-      MarkNode,
-    ],
-  }
-
-  return (
-    <div className={cn("text-text relative text-sm", className)}>
-      <LexicalComposer initialConfig={initialConfig}>
-        <RichTextPlugin
-          contentEditable={
-            <ContentEditable className="focus:outline-none" style={{ outline: "none" }} />
-          }
-          ErrorBoundary={LexicalErrorBoundary}
-          placeholder={null}
-        />
-      </LexicalComposer>
-    </div>
-  )
-}
+    return (
+      <div className={cn("text-text relative text-sm", className)}>
+        <LexicalComposer initialConfig={initialConfig}>
+          <RichTextPlugin
+            contentEditable={
+              <ContentEditable className="focus:outline-none" style={{ outline: "none" }} />
+            }
+            ErrorBoundary={LexicalErrorBoundary}
+            placeholder={null}
+          />
+        </LexicalComposer>
+      </div>
+    )
+  },
+)
--- a/apps/desktop/layer/renderer/src/modules/ai/chat/components/message/EditableMessage.tsx
+++ b/apps/desktop/layer/renderer/src/modules/ai/chat/components/message/EditableMessage.tsx
@ -1,21 +1,29 @@
-import { useInputComposition } from "@follow/hooks"
+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 { useCallback, useEffect, useRef, useState } from "react"
+import type { BizUIMessage } from "@folo-services/ai-tools"
+import { isEqual } from "es-toolkit"
+import type { EditorState, LexicalEditor, SerializedEditorState } from "lexical"
+import { $getRoot } from "lexical"
+import { useCallback, useEffect, useMemo, useRef, useState } from "react"

 import { useChatStatus } from "~/modules/ai/chat/__internal__/hooks"
 import { useEditingMessageId, useSetEditingMessageId } from "~/modules/ai/chat/atoms/session"

 interface EditableMessageProps {
  messageId: string
-  initialContent: string
-  onSave: (content: string) => void
+  parts: BizUIMessage["parts"]
+  onSave: (content: SerializedEditorState, editor: LexicalEditor) => void
  onCancel: () => void
  className?: string
 }

 export const EditableMessage = ({
  messageId,
-  initialContent,
+  parts,
  onSave,
  onCancel,
  className,
@ -23,72 +31,69 @@ export const EditableMessage = ({
  const status = useChatStatus()
  const editingMessageId = useEditingMessageId()
  const setEditingMessageId = useSetEditingMessageId()
-  const [content, setContent] = useState(initialContent)
  const [isEmpty, setIsEmpty] = useState(false)
-  const textareaRef = useRef<HTMLTextAreaElement>(null)
+  const editorRef = useRef<LexicalRichEditorRef>(null)
+  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])
  const isEditing = editingMessageId === messageId
  const isProcessing = status === "submitted" || status === "streaming"

-  // Auto-resize textarea and maintain minimum height to prevent CLS
+  // Initialize editor with initial content
  useEffect(() => {
-    if (textareaRef.current && isEditing) {
-      const textarea = textareaRef.current
-      textarea.style.height = "auto"
-      const newHeight = Math.max(56, textarea.scrollHeight) // Minimum height of 56px
-      textarea.style.height = `${newHeight}px`
+    if (isEditing && editorRef.current && currentEditor) {
+      // Focus the editor
+      editorRef.current.focus()
    }
-  }, [content, isEditing])
-
-  // Focus on edit start
-  useEffect(() => {
-    if (isEditing && textareaRef.current) {
-      textareaRef.current.focus()
-      textareaRef.current.setSelectionRange(
-        textareaRef.current.value.length,
-        textareaRef.current.value.length,
-      )
-    }
-  }, [isEditing])
+  }, [isEditing, initialEditorState, currentEditor])

  const handleSave = useCallback(() => {
-    if (content.trim() && content.trim() !== initialContent) {
-      onSave(content.trim())
+    if (currentEditor && editorRef.current && !editorRef.current.isEmpty()) {
+      const serializedEditorState = currentEditor.getEditorState().toJSON()
+
+      if (!isEqual(serializedEditorState, initialEditorState.toJSON())) {
+        onSave(serializedEditorState, currentEditor)
+      }
    }
-    setEditingMessageId(null)
-  }, [content, initialContent, onSave, setEditingMessageId])
+  }, [currentEditor, initialEditorState, onSave])

  const handleCancel = useCallback(() => {
-    setContent(initialContent)
    setEditingMessageId(null)
    onCancel()
-  }, [initialContent, onCancel, setEditingMessageId])
+  }, [onCancel, setEditingMessageId])

-  const handleKeyPress = useCallback(
-    (e: React.KeyboardEvent) => {
-      if (e.key === "Enter" && !e.shiftKey) {
-        e.preventDefault()
+  const handleKeyDown = useCallback(
+    (event: KeyboardEvent) => {
+      if (event.key === "Enter" && !event.shiftKey) {
+        event.preventDefault()
        if (!isProcessing) {
          handleSave()
        }
-      } else if (e.key === "Escape") {
-        e.preventDefault()
+        return true
+      } else if (event.key === "Escape") {
+        event.preventDefault()
        handleCancel()
+        return true
      }
+      return false
    },
    [handleSave, handleCancel, isProcessing],
  )

-  const handleChange = useCallback((e: React.ChangeEvent<HTMLTextAreaElement>) => {
-    const newContent = e.target.value
-    setContent(newContent)
-    setIsEmpty(newContent.trim() === "")
+  const handleEditorChange = useCallback((editorState: EditorState, editor: LexicalEditor) => {
+    setCurrentEditor(editor)
+    // Update isEmpty state based on editor content
+    editorState.read(() => {
+      const root = $getRoot()
+      const textContent = root.getTextContent().trim()
+      setIsEmpty(textContent === "")
+    })
  }, [])

-  const inputProps = useInputComposition<HTMLTextAreaElement>({
-    onKeyDown: handleKeyPress,
-  })
-
  if (!isEditing) {
    return null
  }
@ -97,15 +102,14 @@ 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">
-        <textarea
-          ref={textareaRef}
-          value={content}
-          onChange={handleChange}
-          {...inputProps}
+        <LexicalRichEditor
+          ref={editorRef}
          placeholder="Edit your message..."
-          className="scrollbar-none text-text placeholder:text-text-secondary max-h-40 min-h-14 w-full resize-none bg-transparent px-4 py-3 pr-20 text-sm !outline-none transition-all duration-200"
-          rows={1}
-          disabled={isProcessing}
+          initalEditorState={initialEditorState}
+          className="w-full pr-20"
+          onChange={handleEditorChange}
+          onKeyDown={handleKeyDown}
+          namespace="EditableMessageRichEditor"
        />

        {/* Action buttons */}
@ -122,7 +126,7 @@ export const EditableMessage = ({
          <button
            type="button"
            onClick={handleSave}
-            disabled={isProcessing || isEmpty || content.trim() === initialContent}
+            disabled={isProcessing || isEmpty}
            className="text-accent hover:text-accent hover:bg-accent/10 flex size-8 items-center justify-center rounded-lg transition-colors disabled:opacity-50"
            title="Save (Enter)"
          >
--- a/apps/desktop/layer/renderer/src/modules/ai/chat/types/ChatSession.ts
+++ b/apps/desktop/layer/renderer/src/modules/ai/chat/types/ChatSession.ts
@ -1,6 +1,16 @@
+import type { SerializedEditorState } from "lexical"
+
 export interface ChatSession {
  chatId: string
  title?: string
  createdAt: Date
  messageCount: number
 }
+
+export type RichTextPart = {
+  type: "data-rich-text"
+  data: {
+    state: SerializedEditorState
+    text: string
+  }
+}
--- a/packages/internal/components/src/ui/lexical-rich-editor/LexicalRichEditor.tsx
+++ b/packages/internal/components/src/ui/lexical-rich-editor/LexicalRichEditor.tsx
@ -1,10 +1,7 @@
 import { cn } from "@follow/utils"
-import { CodeHighlightNode, CodeNode } from "@lexical/code"
-import { LinkNode } from "@lexical/link"
-import { ListItemNode, ListNode } from "@lexical/list"
-import { MarkNode } from "@lexical/mark"
 import { TRANSFORMERS } from "@lexical/markdown"
 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 { LexicalErrorBoundary } from "@lexical/react/LexicalErrorBoundary"
@ -14,11 +11,11 @@ import { ListPlugin } from "@lexical/react/LexicalListPlugin"
 import { MarkdownShortcutPlugin } from "@lexical/react/LexicalMarkdownShortcutPlugin"
 import { OnChangePlugin } from "@lexical/react/LexicalOnChangePlugin"
 import { RichTextPlugin } from "@lexical/react/LexicalRichTextPlugin"
-import { HeadingNode, QuoteNode } from "@lexical/rich-text"
 import type { EditorState, LexicalEditor } from "lexical"
-import { $getRoot, ParagraphNode, TextNode } from "lexical"
+import { $getRoot } from "lexical"
 import { useImperativeHandle, useRef, useState } from "react"

+import { LexicalRichEditorNodes } from "./nodes"
 import { KeyboardPlugin } from "./plugins"
 import { defaultLexicalTheme } from "./theme"
 import type { LexicalRichEditorProps, LexicalRichEditorRef } from "./types"
@ -26,6 +23,13 @@ import type { LexicalRichEditorProps, LexicalRichEditorRef } from "./types"
 function onError(error: Error) {
  console.error("Lexical Editor Error:", error)
 }
+const defaultEnabledPlugins = {
+  history: true,
+  markdown: true,
+  list: true,
+  link: true,
+  autoFocus: true,
+}

 export const LexicalRichEditor = ({
  ref,
@ -36,44 +40,18 @@ export const LexicalRichEditor = ({
  autoFocus = false,
  namespace = "LexicalRichEditor",
  theme = defaultLexicalTheme,
-  enabledPlugins = {
-    history: true,
-    markdown: true,
-    list: true,
-    link: true,
-    autoFocus: true,
-  },
+  enabledPlugins = defaultEnabledPlugins,
+  initalEditorState,
 }: LexicalRichEditorProps & { ref?: React.RefObject<LexicalRichEditorRef | null> }) => {
  const editorRef = useRef<LexicalEditor | null>(null)
  const [isEmpty, setIsEmpty] = useState(true)

-  const initialConfig = {
+  const initialConfig: InitialConfigType = {
    namespace,
    theme,
    onError,
-    nodes: [
-      // Core nodes
-      ParagraphNode,
-      TextNode,
-
-      // Rich text nodes
-      HeadingNode, // For HEADING transformer
-      QuoteNode, // For QUOTE transformer
-
-      // List nodes
-      ListNode, // For UNORDERED_LIST, ORDERED_LIST transformers
-      ListItemNode,
-
-      // Code nodes
-      CodeNode, // For CODE transformer (multiline)
-      CodeHighlightNode, // For code syntax highlighting
-
-      // Link nodes
-      LinkNode, // For LINK transformer
-
-      // Text format nodes
-      MarkNode, // For HIGHLIGHT transformer
-    ],
+    nodes: LexicalRichEditorNodes,
+    editorState: initalEditorState,
  }

  useImperativeHandle(ref, () => ({
--- a/packages/internal/components/src/ui/lexical-rich-editor/editor.tsx
+++ b/packages/internal/components/src/ui/lexical-rich-editor/editor.tsx
@ -0,0 +1,22 @@
+import type { CreateEditorArgs } from "lexical"
+import { createEditor } from "lexical"
+
+import { LexicalRichEditorNodes } from "./nodes"
+import { defaultLexicalTheme } from "./theme"
+
+export const createLexicalEditor = (options: CreateEditorArgs) => {
+  const editor = createEditor({
+    theme: defaultLexicalTheme,
+    nodes: LexicalRichEditorNodes,
+    ...options,
+  })
+  return editor
+}
+
+export const createDefaultLexicalEditor = () => {
+  return createLexicalEditor({
+    namespace: "LexicalRichEditor",
+    theme: defaultLexicalTheme,
+    nodes: LexicalRichEditorNodes,
+  })
+}
--- a/packages/internal/components/src/ui/lexical-rich-editor/index.tsx
+++ b/packages/internal/components/src/ui/lexical-rich-editor/index.tsx
@ -1,4 +1,6 @@
+export { createDefaultLexicalEditor, createLexicalEditor } from "./editor"
 export { LexicalRichEditor } from "./LexicalRichEditor"
+export { LexicalRichEditorNodes } from "./nodes"
 export { KeyboardPlugin } from "./plugins"
 export { defaultLexicalTheme } from "./theme"
 export type { LexicalRichEditorProps, LexicalRichEditorRef } from "./types"
--- a/packages/internal/components/src/ui/lexical-rich-editor/nodes.ts
+++ b/packages/internal/components/src/ui/lexical-rich-editor/nodes.ts
@ -0,0 +1,30 @@
+import { CodeHighlightNode, CodeNode } from "@lexical/code"
+import { LinkNode } from "@lexical/link"
+import { ListItemNode, ListNode } from "@lexical/list"
+import { MarkNode } from "@lexical/mark"
+import { HeadingNode, QuoteNode } from "@lexical/rich-text"
+import { ParagraphNode, TextNode } from "lexical"
+
+export const LexicalRichEditorNodes = [
+  // Core nodes
+  ParagraphNode,
+  TextNode,
+
+  // Rich text nodes
+  HeadingNode, // For HEADING transformer
+  QuoteNode, // For QUOTE transformer
+
+  // List nodes
+  ListNode, // For UNORDERED_LIST, ORDERED_LIST transformers
+  ListItemNode,
+
+  // Code nodes
+  CodeNode, // For CODE transformer (multiline)
+  CodeHighlightNode, // For code syntax highlighting
+
+  // Link nodes
+  LinkNode, // For LINK transformer
+
+  // Text format nodes
+  MarkNode, // For HIGHLIGHT transformer
+]
--- a/packages/internal/components/src/ui/lexical-rich-editor/types.ts
+++ b/packages/internal/components/src/ui/lexical-rich-editor/types.ts
@ -22,4 +22,5 @@ export interface LexicalRichEditorProps {
    link?: boolean
    autoFocus?: boolean
  }
+  initalEditorState?: EditorState
 }