fix(tui): make folded assistant messages expandable

.changeset/expand-folded-assistant-messages.md

@@ -0,0 +1,5 @@
+---
+"@moonshot-ai/kimi-code": patch
+---
+
+Let Ctrl+O reveal assistant replies folded into transcript summaries.

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

@@ -44,6 +44,10 @@ export class AssistantMessageComponent implements Component {
     this.markRenderDirty();
   }
 
+  getContent(): string {
+    return this.lastText;
+  }
+
   updateContent(text: string, opts?: AssistantMarkdownOptions): void {
     const displayText = text.trim();
     const transient = opts?.transient === true;

apps/kimi-code/src/tui/components/messages/step-summary.ts

@@ -1,17 +1,22 @@
 import type { Component } from '@moonshot-ai/pi-tui';
 
+import { AssistantMessageComponent } from '#/tui/components/messages/assistant-message';
 import { currentTheme } from '#/tui/theme';
 
 /**
  * A collapsed summary of older content within a turn. Accumulates counts of
- * merged steps (thinking blocks and tool calls) and folded assistant messages,
- * rendering them as a single muted line, e.g.
+ * merged steps (thinking blocks and tool calls) and folded assistant messages.
+ * The count stays as a single muted line until Ctrl+O reveals lightweight
+ * snapshots of the folded assistant messages, e.g.
  * `… thinking 5 times, call 50 tools, 12 messages`.
  */
 export class StepSummaryComponent implements Component {
   private thinking = 0;
   private tool = 0;
   private message = 0;
+  private foldedMessageContent: string[] = [];
+  private expandedMessageRenderCache: { width: number; lines: string[] } | undefined;
+  private expanded = false;
 
   get isEmpty(): boolean {
     return this.thinking === 0 && this.tool === 0 && this.message === 0;
@@ -23,14 +28,44 @@ export class StepSummaryComponent implements Component {
     this.message += message;
   }
 
-  invalidate(): void {}
+  addFoldedMessages(messages: readonly AssistantMessageComponent[]): void {
+    for (const message of messages) {
+      const content = message.getContent();
+      if (content.length > 0) this.foldedMessageContent.push(content);
+    }
+    this.expandedMessageRenderCache = undefined;
+  }
+
+  setExpanded(expanded: boolean): void {
+    this.expanded = expanded;
+    if (!expanded) this.expandedMessageRenderCache = undefined;
+  }
 
-  render(_width: number): string[] {
+  invalidate(): void {
+    this.expandedMessageRenderCache = undefined;
+  }
+
+  render(width: number): string[] {
     const parts: string[] = [];
     if (this.thinking > 0) parts.push(`thinking ${this.thinking} times`);
     if (this.tool > 0) parts.push(`call ${this.tool} tools`);
     if (this.message > 0) parts.push(`${this.message} messages`);
     if (parts.length === 0) return [];
-    return [currentTheme.dim(`\u2026 ${parts.join(', ')}`)];
+    const lines = [currentTheme.dim(`\u2026 ${parts.join(', ')}`)];
+    if (!this.expanded || this.foldedMessageContent.length === 0) return lines;
+
+    if (this.expandedMessageRenderCache?.width !== width) {
+      // Reuse one Markdown renderer so a very long turn does not rebuild an
+      // equally large component tree when its folded messages are expanded.
+      const message = new AssistantMessageComponent();
+      const expandedLines: string[] = [];
+      for (const content of this.foldedMessageContent) {
+        message.updateContent(content);
+        expandedLines.push(...message.render(width));
+      }
+      this.expandedMessageRenderCache = { width, lines: expandedLines };
+    }
+    lines.push(...this.expandedMessageRenderCache.lines);
+    return lines;
   }
 }

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

@@ -2798,12 +2798,15 @@ export class KimiTUI {
   appendTranscriptEntry(entry: TranscriptEntry): void {
     this.state.transcriptEntries.push(entry);
     const component = this.createTranscriptComponent(entry);
+    let advancesTurn = false;
     if (component) {
       markTranscriptComponent(component, entry);
+      advancesTurn = this.isTurnBoundaryComponent(component);
       this.state.transcriptContainer.addChild(component);
     }
     const trimmed = this.trimTranscriptWindow();
     const merged = this.mergeCurrentTurnSteps();
+    if (advancesTurn) this.applyToolOutputExpansionState();
     if (component || trimmed || merged) {
       this.state.ui.requestRender();
     }
@@ -3047,6 +3050,9 @@ export class KimiTUI {
       ...stepIndices.slice(0, stepMergeCount),
       ...assistantIndices.slice(0, assistantMergeCount),
     ];
+    const foldedAssistants = assistantIndices
+      .slice(0, assistantMergeCount)
+      .map((idx) => children[idx] as AssistantMessageComponent);
 
     let thinkingCount = 0;
     let toolCount = 0;
@@ -3065,6 +3071,8 @@ export class KimiTUI {
       summary = new StepSummaryComponent();
       summary.addCounts(thinkingCount, toolCount, assistantMergeCount);
     }
+    summary.addFoldedMessages(foldedAssistants);
+    summary.setExpanded(this.state.toolOutputExpanded && TRANSCRIPT_EXPAND_TURNS > 0);
 
     // Rebuild children: keep everything except the merged steps, with the summary
     // sitting right after the user message.
@@ -3100,6 +3108,10 @@ export class KimiTUI {
 
     const newChildren: Component[] = [];
     const toDispose: Component[] = [];
+    const expandedTurnStart =
+      TRANSCRIPT_EXPAND_TURNS > 0
+        ? Math.max(0, boundaries.length - TRANSCRIPT_EXPAND_TURNS)
+        : boundaries.length;
     for (let i = 0; i < boundaries[0]!; i++) newChildren.push(children[i]!);
 
     for (let t = 0; t < boundaries.length; t++) {
@@ -3132,6 +3144,9 @@ export class KimiTUI {
           ...stepIndices.slice(0, stepMergeCount),
           ...assistantIndices.slice(0, assistantMergeCount),
         ];
+        const foldedAssistants = assistantIndices
+          .slice(0, assistantMergeCount)
+          .map((idx) => children[idx] as AssistantMessageComponent);
         let thinkingCount = 0;
         let toolCount = 0;
         for (const idx of toMergeIndices) {
@@ -3147,6 +3162,8 @@ export class KimiTUI {
           summary = new StepSummaryComponent();
           summary.addCounts(thinkingCount, toolCount, assistantMergeCount);
         }
+        summary.addFoldedMessages(foldedAssistants);
+        summary.setExpanded(this.state.toolOutputExpanded && t >= expandedTurnStart);
         newChildren.push(summary);
         for (const idx of toMergeIndices) toDispose.push(children[idx]!);
         const toMergeSet = new Set(toMergeIndices);
@@ -3369,6 +3386,14 @@ export class KimiTUI {
 
   toggleToolOutputExpansion(): void {
     this.state.toolOutputExpanded = !this.state.toolOutputExpanded;
+    this.applyToolOutputExpansionState();
+    // Differential render only — no destructive full redraw on expand/collapse.
+    // (When the expanded region reaches above the viewport, the engine's own
+    // fallback may still do a full render; that path is not forced from here.)
+    this.state.ui.requestRender();
+  }
+
+  private applyToolOutputExpansionState(): void {
     const children = this.state.transcriptContainer.children;
 
     // A component is expandable only if it sits at or after the start of the
@@ -3391,10 +3416,6 @@ export class KimiTUI {
       if (!isExpandable(child)) continue;
       child.setExpanded(this.state.toolOutputExpanded && i >= expandCutoff);
     }
-    // Differential render only — no destructive full redraw on expand/collapse.
-    // (When the expanded region reaches above the viewport, the engine's own
-    // fallback may still do a full render; that path is not forced from here.)
-    this.state.ui.requestRender();
   }
 
   toggleTodoPanelExpansion(): void {

apps/kimi-code/test/tui/components/messages/step-summary.test.ts

@@ -1,11 +1,18 @@
 import { describe, expect, it } from 'vitest';
 
+import { AssistantMessageComponent } from '#/tui/components/messages/assistant-message';
 import { StepSummaryComponent } from '#/tui/components/messages/step-summary';
 
 function strip(text: string): string {
   return text.replaceAll(/\u001B\[[0-9;]*m/g, '');
 }
 
+function assistant(text: string): AssistantMessageComponent {
+  const component = new AssistantMessageComponent();
+  component.updateContent(text);
+  return component;
+}
+
 describe('StepSummaryComponent', () => {
   it('renders nothing when empty', () => {
     const component = new StepSummaryComponent();
@@ -32,4 +39,40 @@ describe('StepSummaryComponent', () => {
     expect(out).toContain('call 4 tools');
     expect(out).toContain('8 messages');
   });
+
+  it('reveals folded assistant message snapshots in order when expanded', () => {
+    const component = new StepSummaryComponent();
+    const first = assistant('first hidden reply');
+    component.addCounts(0, 0, 2);
+    component.addFoldedMessages([first, assistant('second hidden reply')]);
+    first.updateContent('mutated after folding');
+
+    const collapsed = strip(component.render(80).join('\n'));
+    expect(collapsed).toContain('2 messages');
+    expect(collapsed).not.toContain('first hidden reply');
+
+    component.setExpanded(true);
+    const expanded = strip(component.render(80).join('\n'));
+    expect(expanded).toContain('first hidden reply');
+    expect(expanded).toContain('second hidden reply');
+    expect(expanded).not.toContain('mutated after folding');
+    expect(expanded.indexOf('first hidden reply')).toBeLessThan(
+      expanded.indexOf('second hidden reply'),
+    );
+
+    component.setExpanded(false);
+    expect(strip(component.render(80).join('\n'))).not.toContain('first hidden reply');
+  });
+
+  it('renders a large set of folded snapshots', () => {
+    const component = new StepSummaryComponent();
+    const messages = Array.from({ length: 100 }, (_, index) => assistant(`hidden reply ${index}`));
+    component.addCounts(0, 0, messages.length);
+    component.addFoldedMessages(messages);
+    component.setExpanded(true);
+
+    const expanded = strip(component.render(80).join('\n'));
+    expect(expanded).toContain('hidden reply 0');
+    expect(expanded).toContain('hidden reply 99');
+  });
 });

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

@@ -30,10 +30,11 @@ import { AssistantMessageComponent } from '#/tui/components/messages/assistant-m
 import { StepSummaryComponent } from '#/tui/components/messages/step-summary';
 import { ToolCallComponent } from '#/tui/components/messages/tool-call';
 import {
-  groupTurns,
+  TRANSCRIPT_EXPAND_TURNS,
   TRANSCRIPT_KEEP_RECENT_ASSISTANT,
   TRANSCRIPT_KEEP_RECENT_ASSISTANT_COMPLETED,
   TRANSCRIPT_KEEP_RECENT_STEPS,
+  groupTurns,
 } from '#/tui/utils/transcript-window';
 import { BtwPanelComponent } from '#/tui/components/panes/btw-panel';
 import { ThinkingComponent } from '#/tui/components/messages/thinking';
@@ -140,6 +141,7 @@ interface MessageDriver {
   closeSession(reason: string): Promise<void>;
   setSession(session: unknown): Promise<void>;
   getCurrentSessionId(): string;
+  toggleToolOutputExpansion(): void;
 }
 
 interface FeedbackDriver extends MessageDriver {
@@ -8276,15 +8278,20 @@ describe('/effort support_efforts override', () => {
 });
 
 describe('transcript step and assistant folding', () => {
-  function driveSteps(driver: MessageDriver, cycles: number): void {
+  function driveSteps(
+    driver: MessageDriver,
+    cycles: number,
+    turnId = 1,
+    messagePrefix = 'msg',
+  ): void {
     for (let i = 0; i < cycles; i++) {
       driver.sessionEventHandler.handleEvent(
         {
           type: 'assistant.delta',
           agentId: 'main',
           sessionId: 'ses-1',
-          turnId: 1,
-          delta: `msg-${i} `,
+          turnId,
+          delta: `${messagePrefix}-${i} `,
         } as Event,
         vi.fn(),
       );
@@ -8293,8 +8300,8 @@ describe('transcript step and assistant folding', () => {
           type: 'tool.call.started',
           agentId: 'main',
           sessionId: 'ses-1',
-          turnId: 1,
-          toolCallId: `call_${i}`,
+          turnId,
+          toolCallId: `call_${turnId}_${i}`,
           name: 'Bash',
           args: { command: 'ls' },
         } as Event,
@@ -8305,8 +8312,8 @@ describe('transcript step and assistant folding', () => {
           type: 'tool.result',
           agentId: 'main',
           sessionId: 'ses-1',
-          turnId: 1,
-          toolCallId: `call_${i}`,
+          turnId,
+          toolCallId: `call_${turnId}_${i}`,
           output: 'ok',
           isError: undefined,
         } as Event,
@@ -8341,6 +8348,14 @@ describe('transcript step and assistant folding', () => {
       (entry) => entry.kind === 'assistant',
     );
     expect(assistantEntries).toHaveLength(cycles);
+
+    const collapsed = stripSgr(renderTranscript(driver));
+    expect(collapsed).not.toContain('msg-0');
+
+    driver.toggleToolOutputExpansion();
+    const expanded = stripSgr(renderTranscript(driver));
+    expect(expanded).toContain('msg-0');
+    expect(expanded.indexOf('msg-0')).toBeLessThan(expanded.indexOf('msg-1'));
   });
 
   it('does not fold a turn within the caps', async () => {
@@ -8392,5 +8407,35 @@ describe('transcript step and assistant folding', () => {
     // The conclusion stays mounted.
     const lastAssistant = assistants.at(-1)!;
     expect(stripSgr(lastAssistant.render(120).join('\n'))).toContain(`msg-${cycles - 1}`);
+
+    expect(stripSgr(renderTranscript(driver))).not.toContain('msg-0');
+    driver.toggleToolOutputExpansion();
+    expect(stripSgr(renderTranscript(driver))).toContain('msg-0');
+  });
+
+  it('collapses expanded summaries after they leave the recent-turn window', async () => {
+    const { driver } = await makeDriver();
+    const cycles = TRANSCRIPT_KEEP_RECENT_ASSISTANT_COMPLETED + 1;
+
+    for (let turnId = 1; turnId <= TRANSCRIPT_EXPAND_TURNS + 1; turnId++) {
+      driver.handleUserInput(`round ${turnId}`);
+      driveSteps(driver, cycles, turnId, `turn-${turnId}-msg`);
+      driver.sessionEventHandler.handleEvent(
+        {
+          type: 'turn.ended',
+          agentId: 'main',
+          sessionId: 'ses-1',
+          turnId,
+          reason: 'completed',
+        } as Event,
+        vi.fn(),
+      );
+      if (turnId === 1) driver.toggleToolOutputExpansion();
+    }
+
+    const transcript = stripSgr(renderTranscript(driver));
+    expect(transcript).not.toContain('turn-1-msg-0');
+    expect(transcript).toContain('turn-2-msg-0');
+    expect(driver.state.toolOutputExpanded).toBe(true);
   });
 });