perf(tui): reuse streaming markdown instances

.changeset/tui-streaming-markdown-reuse.md

@@ -0,0 +1,5 @@
+---
+"@moonshot-ai/kimi-code": patch
+---
+
+Reduce streaming redraw cost for long assistant messages with code blocks.

apps/kimi-code/src/tui/components/messages/assistant-message.ts

@@ -12,9 +12,16 @@ import { STATUS_BULLET } from '#/tui/constant/symbols';
 import { currentTheme } from '#/tui/theme';
 import { createMarkdownTheme } from '#/tui/theme/pi-tui-theme';
 
+type AssistantMarkdownOptions = {
+  transient?: boolean;
+};
+
 export class AssistantMessageComponent implements Component {
   private contentContainer: Container;
+  private markdown: Markdown | undefined;
+  private markdownTransient = false;
   private lastText = '';
+  private lastTransient = false;
   private showBullet: boolean;
 
   constructor(showBullet: boolean = true) {
@@ -26,25 +33,49 @@ export class AssistantMessageComponent implements Component {
     this.showBullet = show;
   }
 
-  updateContent(text: string): void {
-    const displayText = text;
-    if (displayText === this.lastText) return;
+  updateContent(text: string, opts?: AssistantMarkdownOptions): void {
+    const displayText = text.trim();
+    const transient = opts?.transient === true;
+
+    if (displayText === this.lastText && transient === this.lastTransient) return;
+
     this.lastText = displayText;
-    this.contentContainer.clear();
-    if (displayText.trim().length > 0) {
-      this.contentContainer.addChild(new Markdown(displayText.trim(), 0, 0, createMarkdownTheme()));
+    this.lastTransient = transient;
+
+    if (displayText.length === 0) {
+      this.contentContainer.clear();
+      this.markdown = undefined;
+      this.markdownTransient = false;
+      return;
     }
+
+    if (this.markdown === undefined || this.markdownTransient !== transient) {
+      this.contentContainer.clear();
+      this.markdown = new Markdown(displayText, 0, 0, createMarkdownTheme({ transient }));
+      this.markdownTransient = transient;
+      this.contentContainer.addChild(this.markdown);
+      return;
+    }
+
+    this.markdown.setText(displayText);
   }
 
   invalidate(): void {
     // Markdown caches ANSI colour codes keyed on (text, width).  When the
     // theme changes the cached strings contain stale colours, so we rebuild
-    // the Markdown child with the new theme.
+    // the Markdown child with the new theme while preserving transient mode.
     this.contentContainer.clear();
+    this.markdown = undefined;
+
     if (this.lastText.trim().length > 0) {
-      this.contentContainer.addChild(
-        new Markdown(this.lastText.trim(), 0, 0, createMarkdownTheme()),
+      this.markdown = new Markdown(
+        this.lastText.trim(),
+        0,
+        0,
+        createMarkdownTheme({ transient: this.lastTransient }),
       );
+      this.markdownTransient = this.lastTransient;
+      this.contentContainer.addChild(this.markdown);
     }
   }
 

apps/kimi-code/src/tui/controllers/streaming-ui.ts

@@ -600,12 +600,16 @@ export class StreamingUIController {
     const block = this._streamingBlock;
     if (block !== null) {
       block.entry.content = fullText;
-      block.component.updateContent(fullText);
+      block.component.updateContent(fullText, { transient: true });
       this.host.state.ui.requestRender();
     }
   }
 
   onStreamingTextEnd(): void {
+    const block = this._streamingBlock;
+    if (block !== null) {
+      block.component.updateContent(block.entry.content, { transient: false });
+    }
     this._streamingBlock = null;
   }
 

apps/kimi-code/src/tui/theme/pi-tui-theme.ts

@@ -22,7 +22,8 @@ import { currentTheme } from './theme';
 // eslint-disable-next-line no-control-regex -- intentionally matches the ESC byte that opens ANSI SGR sequences.
 const HEADING_HASH_PREFIX = /^((?:\u001B\[[0-9;]*m)*)#{1,6}[ \t]+/;
 
-export function createMarkdownTheme(): MarkdownTheme {
+export function createMarkdownTheme(options?: { transient?: boolean }): MarkdownTheme {
+  const transient = options?.transient === true;
   const stripHash = (text: string): string => text.replace(HEADING_HASH_PREFIX, '$1');
 
   return {
@@ -44,6 +45,8 @@ export function createMarkdownTheme(): MarkdownTheme {
     strikethrough: (text) => chalk.strikethrough(text),
     underline: (text) => chalk.underline(text),
     highlightCode: (code: string, lang?: string) => {
+      if (transient) return code.split('\n');
+
       const normalizedLang = lang?.trim().toLowerCase();
       const language =
         normalizedLang !== undefined && supportsLanguage(normalizedLang) ? normalizedLang : 'text';

apps/kimi-code/test/tui/components/messages/assistant-message.test.ts

@@ -1,12 +1,21 @@
-import { visibleWidth } from '@earendil-works/pi-tui';
-import { describe, expect, it } from 'vitest';
+import { Markdown, visibleWidth } from '@earendil-works/pi-tui';
+import * as cliHighlight from 'cli-highlight';
+import { describe, expect, it, vi } from 'vitest';
 
 import { AssistantMessageComponent } from '#/tui/components/messages/assistant-message';
 import { STATUS_BULLET } from '#/tui/constant/symbols';
 import { createMarkdownTheme } from '#/tui/theme/pi-tui-theme';
 
 import { captureProcessWrite } from '../../../helpers/process';
 
+vi.mock('cli-highlight', async () => {
+  const actual = await vi.importActual<typeof import('cli-highlight')>('cli-highlight');
+  return {
+    ...actual,
+    highlight: vi.fn(actual.highlight),
+  };
+});
+
 function strip(text: string): string {
   return text.replaceAll(/\u001B\[[0-9;]*m/g, '');
 }
@@ -60,4 +69,60 @@ describe('AssistantMessageComponent', () => {
     expect(text).toContain('</hook_result>');
     expect(text).not.toContain('UserPromptSubmit hook');
   });
+
+  it('reuses the same Markdown child across streaming text updates', () => {
+    const component = new AssistantMessageComponent();
+
+    component.updateContent('hello');
+    const first = (component as any).contentContainer.children[0];
+    expect(first).toBeInstanceOf(Markdown);
+
+    component.updateContent('hello world');
+    const second = (component as any).contentContainer.children[0];
+
+    expect(second).toBe(first);
+    expect(strip(component.render(80).join('\n'))).toContain('hello world');
+  });
+
+  it('does not recreate the Markdown child when the text is unchanged', () => {
+    const component = new AssistantMessageComponent();
+
+    component.updateContent('hello');
+    const first = (component as any).contentContainer.children[0];
+    expect(first).toBeInstanceOf(Markdown);
+
+    component.updateContent('hello');
+    const second = (component as any).contentContainer.children[0];
+
+    expect(second).toBe(first);
+  });
+
+  it('rebuilds the Markdown child when transient changes so final render can highlight code', () => {
+    const component = new AssistantMessageComponent();
+    const code = '```ts\nconst x = 1\n```';
+
+    component.updateContent(code, { transient: true });
+    const streaming = (component as any).contentContainer.children[0];
+    expect(streaming).toBeInstanceOf(Markdown);
+
+    component.updateContent(code, { transient: false });
+    const finalized = (component as any).contentContainer.children[0];
+    expect(finalized).toBeInstanceOf(Markdown);
+
+    expect(finalized).not.toBe(streaming);
+  });
+
+  it('skips synchronous syntax highlighting in transient markdown themes', () => {
+    const highlightSpy = vi.mocked(cliHighlight.highlight);
+    highlightSpy.mockClear();
+    const streamingTheme = createMarkdownTheme({ transient: true });
+    const finalTheme = createMarkdownTheme();
+    const code = 'const x = 1';
+
+    expect(streamingTheme.highlightCode?.(code, 'typescript')).toEqual([code]);
+    expect(highlightSpy).not.toHaveBeenCalled();
+
+    finalTheme.highlightCode?.(code, 'typescript');
+    expect(highlightSpy).toHaveBeenCalled();
+  });
 });

apps/kimi-code/test/tui/kimi-tui-message-flow.test.ts

@@ -1326,7 +1326,7 @@ command = "vim"
       await vi.runOnlyPendingTimersAsync();
 
       expect(updateSpy).toHaveBeenCalledTimes(1);
-      expect(updateSpy).toHaveBeenLastCalledWith('abc');
+      expect(updateSpy).toHaveBeenLastCalledWith('abc', { transient: true });
     } finally {
       vi.useRealTimers();
     }