From ada3c6d4471e26e6d51c49e9b5d5f6609819b185 Mon Sep 17 00:00:00 2001
From: LifeJiggy <Bloomtonjovish@gmail.com>
Date: Sat, 30 May 2026 06:08:10 +0100
Subject: [PATCH] fix: invalidate assistant Markdown instances on theme change
 (Codex review)

The pi-tui Markdown component caches rendered output by text+width and
caches the defaultStylePrefix for its entire lifetime. After a theme
switch, existing AssistantMessageComponent instances held stale ANSI
colors because the Markdown child was never rebuilt.

- Add ThemeAwareComponent interface + isThemeAware type guard to
  component-capabilities.ts following the existing Expandable pattern
- Add AssistantMessageComponent.applyTheme() that rebuilds the Markdown
  child with the new markdownTheme and defaultTextStyle
- Wire KimiTUI.applyTheme() to iterate transcript children and call
  applyTheme() on any theme-aware component
---
 .../components/messages/assistant-message.ts   | 18 ++++++++++++++++--
 apps/kimi-code/src/tui/kimi-tui.ts             |  7 ++++++-
 .../src/tui/utils/component-capabilities.ts    | 17 +++++++++++++++++
 3 files changed, 39 insertions(+), 3 deletions(-)

diff --git a/apps/kimi-code/src/tui/components/messages/assistant-message.ts b/apps/kimi-code/src/tui/components/messages/assistant-message.ts
index 1be89b2c..60ce73f3 100644
--- a/apps/kimi-code/src/tui/components/messages/assistant-message.ts
+++ b/apps/kimi-code/src/tui/components/messages/assistant-message.ts
@@ -5,7 +5,7 @@
  * to align after the bullet.
  */
 
-import type { Component, MarkdownTheme } from '@earendil-works/pi-tui';
+import type { Component, DefaultTextStyle, MarkdownTheme } from '@earendil-works/pi-tui';
 import { Container, Markdown, visibleWidth } from '@earendil-works/pi-tui';
 import chalk from 'chalk';
 
@@ -16,12 +16,14 @@ import type { ColorPalette } from '#/tui/theme/colors';
 export class AssistantMessageComponent implements Component {
   private contentContainer: Container;
   private markdownTheme: MarkdownTheme;
+  private defaultTextStyle: DefaultTextStyle;
   private bulletColor: string;
   private lastText = '';
   private showBullet: boolean;
 
   constructor(markdownTheme: MarkdownTheme, colors: ColorPalette, showBullet: boolean = true) {
     this.markdownTheme = markdownTheme;
+    this.defaultTextStyle = { color: (text) => chalk.hex(colors.text)(text) };
     this.bulletColor = colors.roleAssistant;
     this.showBullet = showBullet;
     this.contentContainer = new Container();
@@ -31,13 +33,25 @@ export class AssistantMessageComponent implements Component {
     this.showBullet = show;
   }
 
+  applyTheme(markdownTheme: MarkdownTheme, colors: ColorPalette): void {
+    this.markdownTheme = markdownTheme;
+    this.bulletColor = colors.roleAssistant;
+    this.defaultTextStyle = { color: (text) => chalk.hex(colors.text)(text) };
+    if (this.lastText) {
+      this.contentContainer.clear();
+      this.contentContainer.addChild(
+        new Markdown(this.lastText.trim(), 0, 0, this.markdownTheme, this.defaultTextStyle),
+      );
+    }
+  }
+
   updateContent(text: string): void {
     const displayText = text;
     if (displayText === this.lastText) return;
     this.lastText = displayText;
     this.contentContainer.clear();
     if (displayText.trim().length > 0) {
-      this.contentContainer.addChild(new Markdown(displayText.trim(), 0, 0, this.markdownTheme));
+      this.contentContainer.addChild(new Markdown(displayText.trim(), 0, 0, this.markdownTheme, this.defaultTextStyle));
     }
   }
 
diff --git a/apps/kimi-code/src/tui/kimi-tui.ts b/apps/kimi-code/src/tui/kimi-tui.ts
index db152fe2..ec5b8536 100644
--- a/apps/kimi-code/src/tui/kimi-tui.ts
+++ b/apps/kimi-code/src/tui/kimi-tui.ts
@@ -112,7 +112,7 @@ import {
   type TUIStartupState,
 } from './types';
 import { createTUIState, type TUIState } from './tui-state';
-import { isExpandable, isPlanExpandable } from './utils/component-capabilities';
+import { isExpandable, isPlanExpandable, isThemeAware } from './utils/component-capabilities';
 import { isDeadTerminalError } from './utils/dead-terminal';
 import { formatErrorMessage } from './utils/event-payload';
 import { ImageAttachmentStore, type ImageAttachment } from './utils/image-attachment-store';
@@ -1563,6 +1563,11 @@ export class KimiTUI {
     this.state.theme.styles = nextTheme.styles;
     this.state.theme.markdownTheme = nextTheme.markdownTheme;
     this.setAppState({ theme });
+    for (const child of this.state.transcriptContainer.children) {
+      if (isThemeAware(child)) {
+        child.applyTheme(this.state.theme.markdownTheme, this.state.theme.colors);
+      }
+    }
     this.updateEditorBorderHighlight();
     this.state.ui.requestRender(true);
   }
diff --git a/apps/kimi-code/src/tui/utils/component-capabilities.ts b/apps/kimi-code/src/tui/utils/component-capabilities.ts
index 5f1f0ba9..dd4672f4 100644
--- a/apps/kimi-code/src/tui/utils/component-capabilities.ts
+++ b/apps/kimi-code/src/tui/utils/component-capabilities.ts
@@ -1,3 +1,7 @@
+import type { MarkdownTheme } from '@earendil-works/pi-tui';
+
+import type { ColorPalette } from '#/tui/theme/colors';
+
 export interface Expandable {
   setExpanded(expanded: boolean): void;
 }
@@ -12,6 +16,19 @@ export interface Disposable {
   dispose(): void;
 }
 
+export interface ThemeAwareComponent {
+  applyTheme(markdownTheme: MarkdownTheme, colors: ColorPalette): void;
+}
+
+export function isThemeAware(obj: unknown): obj is ThemeAwareComponent {
+  return (
+    typeof obj === 'object' &&
+    obj !== null &&
+    'applyTheme' in obj &&
+    typeof (obj as ThemeAwareComponent).applyTheme === 'function'
+  );
+}
+
 export function isExpandable(obj: unknown): obj is Expandable {
   return (
     typeof obj === 'object' &&