feat: preview undo ranges before confirmation

.changeset/undo-preview-confirmation.md

@@ -0,0 +1,5 @@
+---
+"@moonshot-ai/kimi-code": minor
+---
+
+Preview `/undo` ranges before confirming and restore the withdrawn prompt to the input box.

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

@@ -125,6 +125,7 @@ export interface SlashCommandHost {
   showSessionPicker(): Promise<void>;
   sendNormalUserInput(text: string): void;
   sendSkillActivation(session: Session, skillName: string, skillArgs: string): void;
+  startUndoPreview(count: number): void;
   readonly skillCommandMap: Map<string, string>;
 
   // Controller refs

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

@@ -5,6 +5,7 @@ import { AgentGroupComponent } from '../components/messages/agent-group';
 import { AssistantMessageComponent } from '../components/messages/assistant-message';
 import { BackgroundAgentStatusComponent } from '../components/messages/background-agent-status';
 import { CronMessageComponent } from '../components/messages/cron-message';
+import { CompactionComponent } from '../components/dialogs/compaction';
 import { ReadGroupComponent } from '../components/messages/read-group';
 import { SkillActivationComponent } from '../components/messages/skill-activation';
 import { ThinkingComponent } from '../components/messages/thinking';
@@ -24,36 +25,43 @@ export async function handleUndoCommand(
   host: SlashCommandHost,
   args: string = '',
 ): Promise<void> {
-  if (host.state.appState.streamingPhase !== 'idle') {
-    host.showError('Cannot undo while streaming — press Esc or Ctrl-C first.');
-    return;
-  }
-
   const count = parseUndoCount(args);
   if (count === undefined) {
     host.showError('Usage: /undo [count], where count is a positive integer.');
     return;
   }
 
+  host.startUndoPreview(count);
+}
+
+export async function commitUndoCommand(
+  host: SlashCommandHost,
+  count: number,
+): Promise<boolean> {
+  if (host.state.appState.streamingPhase !== 'idle') {
+    host.showError('Cannot undo while streaming — press Esc or Ctrl-C first.');
+    return false;
+  }
+
   const session = host.session;
   if (session === undefined) {
     host.showError(NO_ACTIVE_SESSION_MESSAGE);
-    return;
+    return false;
   }
 
   const entries = host.state.transcriptEntries;
   const lastUserIndex = findUndoAnchorEntryIndex(entries, count);
   if (lastUserIndex === undefined) {
     host.showError('Nothing to undo.');
-    return;
+    return false;
   }
 
   try {
     await session.undoHistory(count);
   } catch (error) {
     const message = formatErrorMessage(error);
     host.showError(`Failed to undo: ${message}`);
-    return;
+    return false;
   }
 
   const children = host.state.transcriptContainer.children;
@@ -73,6 +81,7 @@ export async function handleUndoCommand(
   }
 
   host.state.ui.requestRender();
+  return true;
 }
 
 function parseUndoCount(args: string): number | undefined {
@@ -83,7 +92,7 @@ function parseUndoCount(args: string): number | undefined {
   return Number.isSafeInteger(count) ? count : undefined;
 }
 
-function isUndoAnchorEntry(entry: TranscriptEntry): boolean {
+export function isUndoAnchorEntry(entry: TranscriptEntry): boolean {
   return (
     entry.kind === 'user' ||
     (entry.kind === 'skill_activation' && entry.skillTrigger === 'user-slash')
@@ -105,7 +114,7 @@ function findUndoAnchorEntryIndex(
   return undefined;
 }
 
-function isUndoContextEntry(entry: TranscriptEntry): boolean {
+export function isUndoContextEntry(entry: TranscriptEntry): boolean {
   switch (entry.kind) {
     case 'user':
     case 'assistant':
@@ -148,14 +157,19 @@ function removeUndoContextComponents(
   }
 }
 
-function isUndoAnchorComponent(child: Component): boolean {
+export function isUndoAnchorComponent(child: Component): boolean {
   return (
     child instanceof UserMessageComponent ||
     (child instanceof SkillActivationComponent && child.trigger === 'user-slash')
   );
 }
 
-function isUndoContextComponent(child: Component): boolean {
+export function isUndoBoundaryComponent(child: Component): boolean {
+  const entry = getTranscriptComponentEntry(child);
+  return entry?.compactionData !== undefined || child instanceof CompactionComponent;
+}
+
+export function isUndoContextComponent(child: Component): boolean {
   const entry = getTranscriptComponentEntry(child);
   if (entry !== undefined) {
     return isUndoContextEntry(entry);

apps/kimi-code/src/tui/components/editor/custom-editor.ts

@@ -111,6 +111,7 @@ export class CustomEditor extends Editor {
    * through so pi-tui's built-in history navigation runs.
    */
   public onUpArrowEmpty?: () => boolean;
+  public onDownArrowEmpty?: () => boolean;
   public onShiftTab?: () => void;
   /**
    * Called when the user triggers "paste image" (Ctrl-V on Unix,
@@ -320,6 +321,12 @@ export class CustomEditor extends Editor {
       }
     }
 
+    if (matchesKey(normalized, Key.down)) {
+      if (this.getText().length === 0 && this.onDownArrowEmpty) {
+        if (this.onDownArrowEmpty()) return;
+      }
+    }
+
     if (matchesKey(normalized, Key.escape)) {
       if (this.hasAutocompleteActivity()) {
         this.cancelAutocompleteActivity();

apps/kimi-code/src/tui/components/messages/undo-highlight.ts

@@ -0,0 +1,54 @@
+import type { Component } from '@earendil-works/pi-tui';
+import { visibleWidth } from '@earendil-works/pi-tui';
+import chalk from 'chalk';
+
+import type { ColorPalette } from '#/tui/theme/colors';
+
+const HIGHLIGHT_MARK = '┃ ';
+const ELLIPSIS = '…';
+
+interface UndoHighlightOptions {
+  readonly extendBottom?: boolean;
+  readonly maxLines?: number;
+}
+
+export class UndoHighlightComponent implements Component {
+  private readonly childComponents: readonly Component[];
+  private readonly extendBottom: boolean;
+  private readonly maxLines: number | undefined;
+
+  constructor(
+    children: Component | readonly Component[],
+    private readonly colors: ColorPalette,
+    options: UndoHighlightOptions = {},
+  ) {
+    this.childComponents = Array.isArray(children) ? children : [children];
+    this.extendBottom = options.extendBottom ?? false;
+    this.maxLines = options.maxLines;
+  }
+
+  invalidate(): void {
+    for (const child of this.childComponents) {
+      child.invalidate?.();
+    }
+  }
+
+  render(width: number): string[] {
+    const mark = chalk.hex(this.colors.undoHighlight).bold(HIGHLIGHT_MARK);
+    const markWidth = visibleWidth(mark);
+    const childWidth = Math.max(1, width - markWidth);
+    let childLines = this.childComponents.flatMap((child) => child.render(childWidth));
+    const maxLines = this.maxLines;
+    if (maxLines !== undefined && childLines.length > maxLines) {
+      childLines = [
+        ...childLines.slice(0, maxLines),
+        chalk.hex(this.colors.undoHighlight)(ELLIPSIS),
+      ];
+    }
+    const lines = childLines.map((line) => mark + line);
+    if (this.extendBottom) {
+      lines.push(mark);
+    }
+    return lines;
+  }
+}

apps/kimi-code/src/tui/controllers/editor-keyboard.ts

@@ -22,6 +22,11 @@ export interface EditorKeyboardHost {
   cancelInFlight: (() => void) | undefined;
 
   handleUserInput(text: string): void;
+  isUndoPreviewActive(): boolean;
+  moveUndoPreviewUp(): boolean;
+  moveUndoPreviewDown(): boolean;
+  commitUndoPreview(): void;
+  cancelUndoPreview(): boolean;
   steerMessage(session: Session, input: string[]): void;
   recallLastQueued(): string | undefined;
   showError(msg: string): void;
@@ -50,15 +55,29 @@ export class EditorKeyboardController {
     const editor = host.state.editor;
 
     editor.onSubmit = (text: string) => {
+      if (host.isUndoPreviewActive() && text.trim().length === 0) {
+        host.commitUndoPreview();
+        return;
+      }
+      if (host.isUndoPreviewActive()) host.cancelUndoPreview();
       host.handleUserInput(text);
     };
 
     editor.onChange = (text: string) => {
       if (this.pendingExit) this.clearPendingExit();
+      if (host.isUndoPreviewActive() && text.length > 0) {
+        host.cancelUndoPreview();
+      }
       host.updateEditorBorderHighlight(text);
     };
 
     editor.onCtrlC = () => {
+      if (host.isUndoPreviewActive()) {
+        this.clearPendingExit();
+        host.cancelUndoPreview();
+        return;
+      }
+
       if (host.cancelInFlight !== undefined) {
         const cancel = host.cancelInFlight;
         host.cancelInFlight = undefined;
@@ -92,6 +111,12 @@ export class EditorKeyboardController {
     };
 
     editor.onCtrlD = () => {
+      if (host.isUndoPreviewActive()) {
+        this.clearPendingExit();
+        host.cancelUndoPreview();
+        return;
+      }
+
       if (this.pendingExit?.kind === 'ctrl-d') {
         this.clearPendingExit();
         void host.stop();
@@ -102,6 +127,10 @@ export class EditorKeyboardController {
 
     editor.onEscape = () => {
       if (this.pendingExit) this.clearPendingExit();
+      if (host.isUndoPreviewActive()) {
+        host.cancelUndoPreview();
+        return;
+      }
       if (host.state.activeDialog === 'session-picker') {
         host.hideSessionPicker();
         return;
@@ -177,6 +206,10 @@ export class EditorKeyboardController {
     };
 
     editor.onUpArrowEmpty = () => {
+      if (host.isUndoPreviewActive()) {
+        host.moveUndoPreviewUp();
+        return true;
+      }
       if (host.state.appState.streamingPhase === 'idle' && !host.state.appState.isCompacting) return false;
       const recalled = host.recallLastQueued();
       if (recalled !== undefined) {
@@ -188,6 +221,12 @@ export class EditorKeyboardController {
       return false;
     };
 
+    editor.onDownArrowEmpty = () => {
+      if (!host.isUndoPreviewActive()) return false;
+      host.moveUndoPreviewDown();
+      return true;
+    };
+
     editor.onPasteImage = async () => this.handleClipboardImagePaste();
   }
 

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

@@ -10,6 +10,7 @@ import { STREAMING_UI_FLUSH_MS } from '../constant/streaming';
 import { hasDispose } from '../utils/component-capabilities';
 import { appendStreamingArgsPreview, parseStreamingArgs } from '../utils/event-payload';
 import { notifyTerminalOnce } from '../utils/terminal-notification';
+import { markTranscriptComponent } from '../utils/transcript-component-metadata';
 import { nextTranscriptId } from '../utils/transcript-id';
 import type { TodoItem } from '../components/chrome/todo-panel';
 import type {
@@ -546,6 +547,7 @@ export class StreamingUIController {
     );
     this._streamingBlock = { component, entry };
     this.host.pushTranscriptEntry(entry);
+    markTranscriptComponent(component, entry);
     state.transcriptContainer.addChild(component);
     state.ui.requestRender();
   }