From f09fd5aa6989b54949073e3bf3cd0fa66de8d83f Mon Sep 17 00:00:00 2001
From: t8y2 <1156263951@qq.com>
Date: Mon, 18 May 2026 17:22:46 +0800
Subject: [PATCH] perf(ai): cache rendered message segments

---
 .../src/components/editor/AiAssistant.vue     | 45 +--------
 apps/desktop/src/lib/aiMessageRender.ts       | 93 +++++++++++++++++++
 packages/app-tests/aiMessageRender.test.ts    | 42 +++++++++
 3 files changed, 140 insertions(+), 40 deletions(-)
 create mode 100644 apps/desktop/src/lib/aiMessageRender.ts
 create mode 100644 packages/app-tests/aiMessageRender.test.ts

diff --git a/apps/desktop/src/components/editor/AiAssistant.vue b/apps/desktop/src/components/editor/AiAssistant.vue
index 2f68e1565..bb1de9b38 100644
--- a/apps/desktop/src/components/editor/AiAssistant.vue
+++ b/apps/desktop/src/components/editor/AiAssistant.vue
@@ -49,6 +49,7 @@ import { useToast } from "@/composables/useToast";
 import { buildAiContext, runAiStream, type AiAction } from "@/lib/ai";
 import { buildAiAgentPlan } from "@/lib/aiAgentPlan";
 import { buildAiAgentStepItems, type AiAgentStepItem, type AiAgentStepTone } from "@/lib/aiAgentStepPresentation";
+import { createAiMessageRenderer } from "@/lib/aiMessageRender";
 import { Marked } from "marked";
 import {
   aiCancelStream,
@@ -624,44 +625,6 @@ function triggerAction(action: AiAction, instruction?: string) {
 
 defineExpose({ triggerAction });
 
-interface MessageSegment {
-  type: "text" | "code";
-  content: string;
-  lang?: string;
-}
-
-function parseMessage(text: string): MessageSegment[] {
-  const segments: MessageSegment[] = [];
-  const lines = text.split("\n");
-  let i = 0;
-
-  while (i < lines.length) {
-    const fenceMatch = lines[i].match(/^```(sql|mysql|postgresql|sqlite|tsql|clickhouse)?\s*$/i);
-    if (fenceMatch) {
-      const lang = (fenceMatch[1] || "sql").toUpperCase();
-      const codeLines: string[] = [];
-      i++;
-      while (i < lines.length && !/^```\s*$/.test(lines[i])) {
-        codeLines.push(lines[i]);
-        i++;
-      }
-      if (i < lines.length) i++;
-      const content = codeLines.join("\n").trim();
-      if (content) segments.push({ type: "code", lang, content });
-    } else {
-      const textLines: string[] = [];
-      while (i < lines.length && !/^```(sql|mysql|postgresql|sqlite|tsql|clickhouse)?\s*$/i.test(lines[i])) {
-        textLines.push(lines[i]);
-        i++;
-      }
-      const content = textLines.join("\n");
-      if (content.trim()) segments.push({ type: "text", content });
-    }
-  }
-
-  return segments;
-}
-
 const markedInstance = new Marked({
   breaks: true,
   gfm: true,
@@ -675,6 +638,8 @@ const markedInstance = new Marked({
 function formatInlineText(text: string): string {
   return markedInstance.parse(text) as string;
 }
+
+const messageRenderer = createAiMessageRenderer({ markdown: formatInlineText });
 </script>
 
 <template>
@@ -792,9 +757,9 @@ function formatInlineText(text: string): string {
                   <span class="truncate">{{ t(step.labelKey) }}</span>
                 </span>
               </div>
-              <template v-for="(seg, j) in parseMessage(msg.content)" :key="j">
+              <template v-for="(seg, j) in messageRenderer.render(msg.content)" :key="j">
                 <div v-if="seg.type === 'text'" class="ai-markdown whitespace-normal">
-                  <div v-html="formatInlineText(seg.content)" />
+                  <div v-html="seg.html" />
                 </div>
                 <div v-else class="my-2 rounded-md overflow-hidden bg-zinc-900 dark:bg-zinc-900">
                   <div
diff --git a/apps/desktop/src/lib/aiMessageRender.ts b/apps/desktop/src/lib/aiMessageRender.ts
new file mode 100644
index 000000000..e57fe430e
--- /dev/null
+++ b/apps/desktop/src/lib/aiMessageRender.ts
@@ -0,0 +1,93 @@
+export interface AiMessageTextSegment {
+  type: "text";
+  content: string;
+  html: string;
+}
+
+export interface AiMessageCodeSegment {
+  type: "code";
+  content: string;
+  lang: string;
+}
+
+export type AiMessageRenderSegment = AiMessageTextSegment | AiMessageCodeSegment;
+
+interface MessageSegment {
+  type: "text" | "code";
+  content: string;
+  lang?: string;
+}
+
+export interface AiMessageRendererOptions {
+  maxEntries?: number;
+  markdown: (text: string) => string;
+}
+
+const DEFAULT_MAX_ENTRIES = 100;
+
+export function createAiMessageRenderer(options: AiMessageRendererOptions) {
+  const maxEntries = Math.max(1, Math.floor(options.maxEntries ?? DEFAULT_MAX_ENTRIES));
+  const cache = new Map<string, AiMessageRenderSegment[]>();
+
+  function render(content: string): AiMessageRenderSegment[] {
+    const cached = cache.get(content);
+    if (cached) {
+      cache.delete(content);
+      cache.set(content, cached);
+      return cached;
+    }
+
+    const rendered = parseAiMessage(content).map((segment): AiMessageRenderSegment => {
+      if (segment.type === "text") {
+        return { type: "text", content: segment.content, html: options.markdown(segment.content) };
+      }
+      return { type: "code", content: segment.content, lang: segment.lang ?? "SQL" };
+    });
+
+    cache.set(content, rendered);
+    while (cache.size > maxEntries) {
+      const oldestKey = cache.keys().next().value;
+      if (oldestKey === undefined) break;
+      cache.delete(oldestKey);
+    }
+    return rendered;
+  }
+
+  function clear() {
+    cache.clear();
+  }
+
+  return { render, clear };
+}
+
+export function parseAiMessage(text: string): MessageSegment[] {
+  const segments: MessageSegment[] = [];
+  const lines = text.split("\n");
+  let i = 0;
+
+  while (i < lines.length) {
+    const fenceMatch = lines[i].match(/^```(sql|mysql|postgresql|sqlite|tsql|clickhouse)?\s*$/i);
+    if (fenceMatch) {
+      const lang = (fenceMatch[1] || "sql").toUpperCase();
+      const codeLines: string[] = [];
+      i++;
+      while (i < lines.length && !/^```\s*$/.test(lines[i])) {
+        codeLines.push(lines[i]);
+        i++;
+      }
+      if (i < lines.length) i++;
+      const content = codeLines.join("\n").trim();
+      if (content) segments.push({ type: "code", lang, content });
+    } else {
+      const textLines: string[] = [];
+      while (i < lines.length && !/^```(sql|mysql|postgresql|sqlite|tsql|clickhouse)?\s*$/i.test(lines[i])) {
+        textLines.push(lines[i]);
+        i++;
+      }
+      const content = textLines.join("\n");
+      if (content.trim()) segments.push({ type: "text", content });
+    }
+  }
+
+  return segments;
+}
diff --git a/packages/app-tests/aiMessageRender.test.ts b/packages/app-tests/aiMessageRender.test.ts
new file mode 100644
index 000000000..6628ebf6f
--- /dev/null
+++ b/packages/app-tests/aiMessageRender.test.ts
@@ -0,0 +1,42 @@
+import assert from "node:assert/strict";
+import test from "node:test";
+import { createAiMessageRenderer } from "../../apps/desktop/src/lib/aiMessageRender.ts";
+
+test("reuses rendered AI message segments for unchanged content", () => {
+  let markdownCalls = 0;
+  const renderer = createAiMessageRenderer({
+    markdown: (text) => {
+      markdownCalls++;
+      return `<p>${text}</p>`;
+    },
+  });
+
+  const first = renderer.render("hello **dbx**\n```sql\nSELECT 1\n```");
+  const second = renderer.render("hello **dbx**\n```sql\nSELECT 1\n```");
+
+  assert.equal(markdownCalls, 1);
+  assert.strictEqual(second, first);
+  assert.deepEqual(second, [
+    { type: "text", content: "hello **dbx**", html: "<p>hello **dbx**</p>" },
+    { type: "code", content: "SELECT 1", lang: "SQL" },
+  ]);
+});
+
+test("evicts older rendered AI message cache entries", () => {
+  let markdownCalls = 0;
+  const renderer = createAiMessageRenderer({
+    maxEntries: 2,
+    markdown: (text) => {
+      markdownCalls++;
+      return text;
+    },
+  });
+
+  renderer.render("one");
+  renderer.render("two");
+  renderer.render("one");
+  renderer.render("three");
+  renderer.render("two");
+
+  assert.equal(markdownCalls, 4);
+});