Fix bugs where Rewind and Resume showed Ugly and 100X too verbose content. (#17940)

2026-02-01 22:48:03 +00:00 · 2026-01-30 10:09:27 -08:00
parent f14d0c6a17
commit bb6a336ca9
16 changed files with 212 additions and 20 deletions
--- a/packages/cli/src/nonInteractiveCli.test.ts
+++ b/packages/cli/src/nonInteractiveCli.test.ts
@@ -258,6 +258,9 @@ describe('runNonInteractive', () => {
      [{ text: 'Test input' }],
      expect.any(AbortSignal),
      'prompt-id-1',
+      undefined,
+      false,
+      'Test input',
    );
    expect(getWrittenOutput()).toBe('Hello World\n');
    // Note: Telemetry shutdown is now handled in runExitCleanup() in cleanup.ts
@@ -374,6 +377,9 @@ describe('runNonInteractive', () => {
      [{ text: 'Tool response' }],
      expect.any(AbortSignal),
      'prompt-id-2',
+      undefined,
+      false,
+      undefined,
    );
    expect(getWrittenOutput()).toBe('Final answer\n');
  });
@@ -531,6 +537,9 @@ describe('runNonInteractive', () => {
      ],
      expect.any(AbortSignal),
      'prompt-id-3',
+      undefined,
+      false,
+      undefined,
    );
    expect(getWrittenOutput()).toBe('Sorry, let me try again.\n');
  });
@@ -670,6 +679,9 @@ describe('runNonInteractive', () => {
      processedParts,
      expect.any(AbortSignal),
      'prompt-id-7',
+      undefined,
+      false,
+      rawInput,
    );

    // 6. Assert the final output is correct
@@ -703,6 +715,9 @@ describe('runNonInteractive', () => {
      [{ text: 'Test input' }],
      expect.any(AbortSignal),
      'prompt-id-1',
+      undefined,
+      false,
+      'Test input',
    );
    expect(processStdoutSpy).toHaveBeenCalledWith(
      JSON.stringify(
@@ -833,6 +848,9 @@ describe('runNonInteractive', () => {
      [{ text: 'Empty response test' }],
      expect.any(AbortSignal),
      'prompt-id-empty',
+      undefined,
+      false,
+      'Empty response test',
    );

    // This should output JSON with empty response but include stats
@@ -967,6 +985,9 @@ describe('runNonInteractive', () => {
      [{ text: 'Prompt from command' }],
      expect.any(AbortSignal),
      'prompt-id-slash',
+      undefined,
+      false,
+      '/testcommand',
    );

    expect(getWrittenOutput()).toBe('Response from command\n');
@@ -1010,6 +1031,9 @@ describe('runNonInteractive', () => {
      [{ text: 'Slash command output' }],
      expect.any(AbortSignal),
      'prompt-id-slash',
+      undefined,
+      false,
+      '/help',
    );
    expect(getWrittenOutput()).toBe('Response to slash command\n');
    handleSlashCommandSpy.mockRestore();
@@ -1184,6 +1208,9 @@ describe('runNonInteractive', () => {
      [{ text: '/unknowncommand' }],
      expect.any(AbortSignal),
      'prompt-id-unknown',
+      undefined,
+      false,
+      '/unknowncommand',
    );

    expect(getWrittenOutput()).toBe('Response to unknown\n');
--- a/packages/cli/src/nonInteractiveCli.ts
+++ b/packages/cli/src/nonInteractiveCli.ts
@@ -301,6 +301,9 @@ export async function runNonInteractive({
          currentMessages[0]?.parts || [],
          abortController.signal,
          prompt_id,
+          undefined,
+          false,
+          turnCount === 1 ? input : undefined,
        );

        let responseText = '';
--- a/packages/cli/src/ui/components/RewindViewer.test.tsx
+++ b/packages/cli/src/ui/components/RewindViewer.test.tsx
@@ -254,6 +254,7 @@ describe('RewindViewer', () => {
      {
        description: 'removes reference markers',
        prompt: `some command @file\n--- Content from referenced files ---\nContent from file:\nblah blah\n--- End of content ---`,
+        expected: 'some command @file',
      },
      {
        description: 'strips expanded MCP resource content',
@@ -263,10 +264,23 @@ describe('RewindViewer', () => {
          '\nContent from @server3:mcp://demo-resource:\n' +
          'This is the content of the demo resource.\n' +
          `--- End of content ---`,
+        expected: 'read @server3:mcp://demo-resource hello',
      },
-    ])('$description', async ({ prompt }) => {
+      {
+        description: 'uses displayContent if present and does not strip',
+        prompt: `raw content with markers\n--- Content from referenced files ---\nblah\n--- End of content ---`,
+        displayContent: 'clean display content',
+        expected: 'clean display content',
+      },
+    ])('$description', async ({ prompt, displayContent, expected }) => {
      const conversation = createConversation([
-        { type: 'user', content: prompt, id: '1', timestamp: '1' },
+        {
+          type: 'user',
+          content: prompt,
+          displayContent,
+          id: '1',
+          timestamp: '1',
+        },
      ]);
      const onRewind = vi.fn();
      const { lastFrame, stdin } = renderWithProviders(
@@ -289,6 +303,15 @@ describe('RewindViewer', () => {
      await waitFor(() => {
        expect(lastFrame()).toContain('Confirm Rewind');
      });
+
+      // Confirm
+      act(() => {
+        stdin.write('\r');
+      });
+
+      await waitFor(() => {
+        expect(onRewind).toHaveBeenCalledWith('1', expected, expect.anything());
+      });
    });
  });

--- a/packages/cli/src/ui/components/RewindViewer.tsx
+++ b/packages/cli/src/ui/components/RewindViewer.tsx
@@ -35,6 +35,14 @@ interface RewindViewerProps {

 const MAX_LINES_PER_BOX = 2;

+const getCleanedRewindText = (userPrompt: MessageRecord): string => {
+  const contentToUse = userPrompt.displayContent || userPrompt.content;
+  const originalUserText = contentToUse ? partToString(contentToUse) : '';
+  return userPrompt.displayContent
+    ? originalUserText
+    : stripReferenceContent(originalUserText);
+};
+
 export const RewindViewer: React.FC<RewindViewerProps> = ({
  conversation,
  onExit,
@@ -162,10 +170,7 @@ export const RewindViewer: React.FC<RewindViewerProps> = ({
              (m) => m.id === selectedMessageId,
            );
            if (userPrompt) {
-              const originalUserText = userPrompt.content
-                ? partToString(userPrompt.content)
-                : '';
-              const cleanedText = stripReferenceContent(originalUserText);
+              const cleanedText = getCleanedRewindText(userPrompt);
              setIsRewinding(true);
              await onRewind(selectedMessageId, cleanedText, outcome);
            }
@@ -224,7 +229,9 @@ export const RewindViewer: React.FC<RewindViewerProps> = ({
                      isSelected ? theme.status.success : theme.text.primary
                    }
                  >
-                    {partToString(userPrompt.content)}
+                    {partToString(
+                      userPrompt.displayContent || userPrompt.content,
+                    )}
                  </Text>
                  <Text color={theme.text.secondary}>
                    Cancel rewind and stay here
@@ -235,10 +242,7 @@ export const RewindViewer: React.FC<RewindViewerProps> = ({

            const stats = getStats(userPrompt);
            const firstFileName = stats?.details?.at(0)?.fileName;
-            const originalUserText = userPrompt.content
-              ? partToString(userPrompt.content)
-              : '';
-            const cleanedText = stripReferenceContent(originalUserText);
+            const cleanedText = getCleanedRewindText(userPrompt);

            return (
              <Box flexDirection="column" marginBottom={1}>
--- a/packages/cli/src/ui/components/snapshots/RewindViewer.test.tsx.snap
+++ b/packages/cli/src/ui/components/snapshots/RewindViewer.test.tsx.snap
@@ -34,6 +34,23 @@ exports[`RewindViewer > Content Filtering > 'strips expanded MCP resource conten
 ╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯"
 `;

+exports[`RewindViewer > Content Filtering > 'uses displayContent if present and do…' 1`] = `
+"╭──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╮
+│                                                                                                                      │
+│ > Rewind                                                                                                             │
+│                                                                                                                      │
+│   clean display content                                                                                              │
+│   No files have been changed                                                                                         │
+│                                                                                                                      │
+│ ● Stay at current position                                                                                           │
+│   Cancel rewind and stay here                                                                                        │
+│                                                                                                                      │
+│                                                                                                                      │
+│ (Use Enter to select a message, Esc to close, Right/Left to expand/collapse)                                         │
+│                                                                                                                      │
+╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯"
+`;
+
 exports[`RewindViewer > Interaction Selection > 'cancels on Escape' > confirmation-dialog 1`] = `
 "╭──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╮
 │                                                                                                                      │
--- a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx
+++ b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx
@@ -655,6 +655,9 @@ describe('useGeminiStream', () => {
      expectedMergedResponse,
      expect.any(AbortSignal),
      'prompt-id-2',
+      undefined,
+      false,
+      expectedMergedResponse,
    );
  });

@@ -1057,6 +1060,9 @@ describe('useGeminiStream', () => {
        toolCallResponseParts,
        expect.any(AbortSignal),
        'prompt-id-4',
+        undefined,
+        false,
+        toolCallResponseParts,
      );
    });

@@ -1498,6 +1504,9 @@ describe('useGeminiStream', () => {
          'This is the actual prompt from the command file.',
          expect.any(AbortSignal),
          expect.any(String),
+          undefined,
+          false,
+          '/my-custom-command',
        );

        expect(mockScheduleToolCalls).not.toHaveBeenCalled();
@@ -1524,6 +1533,9 @@ describe('useGeminiStream', () => {
          '',
          expect.any(AbortSignal),
          expect.any(String),
+          undefined,
+          false,
+          '/emptycmd',
        );
      });
    });
@@ -1542,6 +1554,9 @@ describe('useGeminiStream', () => {
          '// This is a line comment',
          expect.any(AbortSignal),
          expect.any(String),
+          undefined,
+          false,
+          '// This is a line comment',
        );
      });
    });
@@ -1560,6 +1575,9 @@ describe('useGeminiStream', () => {
          '/* This is a block comment */',
          expect.any(AbortSignal),
          expect.any(String),
+          undefined,
+          false,
+          '/* This is a block comment */',
        );
      });
    });
@@ -2392,6 +2410,9 @@ describe('useGeminiStream', () => {
      processedQueryParts, // Argument 1: The parts array directly
      expect.any(AbortSignal), // Argument 2: An AbortSignal
      expect.any(String), // Argument 3: The prompt_id string
+      undefined,
+      false,
+      rawQuery,
    );
  });

@@ -2931,6 +2952,9 @@ describe('useGeminiStream', () => {
          'test query',
          expect.any(AbortSignal),
          expect.any(String),
+          undefined,
+          false,
+          'test query',
        );
      });
    });
@@ -3078,6 +3102,9 @@ describe('useGeminiStream', () => {
          'second query',
          expect.any(AbortSignal),
          expect.any(String),
+          undefined,
+          false,
+          'second query',
        );
      });
    });
--- a/packages/cli/src/ui/hooks/useGeminiStream.ts
+++ b/packages/cli/src/ui/hooks/useGeminiStream.ts
@@ -1255,6 +1255,9 @@ export const useGeminiStream = (
                queryToSend,
                abortSignal,
                prompt_id!,
+                undefined,
+                false,
+                query,
              );
              const processingStatus = await processGeminiStreamEvents(
                stream,
--- a/packages/cli/src/ui/hooks/useSessionBrowser.test.ts
+++ b/packages/cli/src/ui/hooks/useSessionBrowser.test.ts
@@ -178,6 +178,30 @@ describe('convertSessionToHistoryFormats', () => {
    });
  });

+  it('should prioritize displayContent for UI history but use content for client history', () => {
+    const messages: MessageRecord[] = [
+      {
+        type: 'user',
+        content: [{ text: 'Expanded content' }],
+        displayContent: [{ text: 'User input' }],
+      } as MessageRecord,
+    ];
+
+    const result = convertSessionToHistoryFormats(messages);
+
+    expect(result.uiHistory).toHaveLength(1);
+    expect(result.uiHistory[0]).toMatchObject({
+      type: 'user',
+      text: 'User input',
+    });
+
+    expect(result.clientHistory).toHaveLength(1);
+    expect(result.clientHistory[0]).toEqual({
+      role: 'user',
+      parts: [{ text: 'Expanded content' }],
+    });
+  });
+
  it('should filter out slash commands from client history but keep in UI', () => {
    const messages: MessageRecord[] = [
      { type: 'user', content: '/help' } as MessageRecord,
--- a/packages/cli/src/ui/hooks/useSessionBrowser.ts
+++ b/packages/cli/src/ui/hooks/useSessionBrowser.ts
@@ -15,6 +15,7 @@ import type {
 } from '@google/gemini-cli-core';
 import type { Part } from '@google/genai';
 import { partListUnionToString, coreEvents } from '@google/gemini-cli-core';
+import { checkExhaustive } from '../../utils/checks.js';
 import type { SessionInfo } from '../../utils/sessionUtils.js';
 import { MessageType, ToolCallStatus } from '../types.js';

@@ -125,8 +126,13 @@ export function convertSessionToHistoryFormats(

  for (const msg of messages) {
    // Add the message only if it has content
+    const displayContentString = msg.displayContent
+      ? partListUnionToString(msg.displayContent)
+      : undefined;
    const contentString = partListUnionToString(msg.content);
-    if (msg.content && contentString.trim()) {
+    const uiText = displayContentString || contentString;
+
+    if (uiText.trim()) {
      let messageType: MessageType;
      switch (msg.type) {
        case 'user':
@@ -141,14 +147,18 @@ export function convertSessionToHistoryFormats(
        case 'warning':
          messageType = MessageType.WARNING;
          break;
+        case 'gemini':
+          messageType = MessageType.GEMINI;
+          break;
        default:
+          checkExhaustive(msg);
          messageType = MessageType.GEMINI;
          break;
      }

      uiHistory.push({
        type: messageType,
-        text: contentString,
+        text: uiText,
      });
    }

@@ -199,7 +209,9 @@ export function convertSessionToHistoryFormats(
      // Add regular user message
      clientHistory.push({
        role: 'user',
-        parts: [{ text: contentString }],
+        parts: Array.isArray(msg.content)
+          ? (msg.content as Part[])
+          : [{ text: contentString }],
      });
    } else if (msg.type === 'gemini') {
      // Handle Gemini messages with potential tool calls
--- a/packages/core/src/core/client.test.ts
+++ b/packages/core/src/core/client.test.ts
@@ -890,6 +890,7 @@ ${JSON.stringify(
        { model: 'default-routed-model' },
        initialRequest,
        expect.any(AbortSignal),
+        undefined,
      );
    });

@@ -1707,6 +1708,7 @@ ${JSON.stringify(
          { model: 'routed-model' },
          [{ text: 'Hi' }],
          expect.any(AbortSignal),
+          undefined,
        );
      });

@@ -1724,6 +1726,7 @@ ${JSON.stringify(
          { model: 'routed-model' },
          [{ text: 'Hi' }],
          expect.any(AbortSignal),
+          undefined,
        );

        // Second turn
@@ -1741,6 +1744,7 @@ ${JSON.stringify(
          { model: 'routed-model' },
          [{ text: 'Continue' }],
          expect.any(AbortSignal),
+          undefined,
        );
      });

@@ -1758,6 +1762,7 @@ ${JSON.stringify(
          { model: 'routed-model' },
          [{ text: 'Hi' }],
          expect.any(AbortSignal),
+          undefined,
        );

        // New prompt
@@ -1779,6 +1784,7 @@ ${JSON.stringify(
          { model: 'new-routed-model' },
          [{ text: 'A new topic' }],
          expect.any(AbortSignal),
+          undefined,
        );
      });

@@ -1806,6 +1812,7 @@ ${JSON.stringify(
          { model: 'original-model' },
          [{ text: 'Hi' }],
          expect.any(AbortSignal),
+          undefined,
        );

        mockRouterService.route.mockResolvedValue({
@@ -1828,6 +1835,7 @@ ${JSON.stringify(
          { model: 'fallback-model' },
          [{ text: 'Continue' }],
          expect.any(AbortSignal),
+          undefined,
        );
      });
    });
@@ -1912,6 +1920,7 @@ ${JSON.stringify(
        { model: 'default-routed-model' },
        initialRequest,
        expect.any(AbortSignal),
+        undefined,
      );

      // Second call with "Please continue."
@@ -1920,6 +1929,7 @@ ${JSON.stringify(
        { model: 'default-routed-model' },
        [{ text: 'System: Please continue.' }],
        expect.any(AbortSignal),
+        undefined,
      );
    });

@@ -2332,6 +2342,7 @@ ${JSON.stringify(
          expect.objectContaining({ model: 'model-a' }),
          expect.anything(),
          expect.anything(),
+          undefined,
        );
      });

@@ -3183,6 +3194,7 @@ ${JSON.stringify(
          expect.anything(),
          [{ text: 'Please explain' }],
          expect.anything(),
+          undefined,
        );
      });

--- a/packages/core/src/core/client.ts
+++ b/packages/core/src/core/client.ts
@@ -532,6 +532,7 @@ export class GeminiClient {
    prompt_id: string,
    boundedTurns: number,
    isInvalidStreamRetry: boolean,
+    displayContent?: PartListUnion,
  ): AsyncGenerator<ServerGeminiStreamEvent, Turn> {
    // Re-initialize turn (it was empty before if in loop, or new instance)
    let turn = new Turn(this.getChat(), prompt_id);
@@ -647,7 +648,12 @@ export class GeminiClient {
      yield { type: GeminiEventType.ModelInfo, value: modelToUse };
    }
    this.currentSequenceModel = modelToUse;
-    const resultStream = turn.run(modelConfigKey, request, linkedSignal);
+    const resultStream = turn.run(
+      modelConfigKey,
+      request,
+      linkedSignal,
+      displayContent,
+    );
    let isError = false;
    let isInvalidStream = false;

@@ -708,6 +714,7 @@ export class GeminiClient {
          prompt_id,
          boundedTurns - 1,
          true,
+          displayContent,
        );
        return turn;
      }
@@ -739,7 +746,8 @@ export class GeminiClient {
            signal,
            prompt_id,
            boundedTurns - 1,
-            // isInvalidStreamRetry is false
+            false, // isInvalidStreamRetry is false
+            displayContent,
          );
          return turn;
        }
@@ -754,6 +762,7 @@ export class GeminiClient {
    prompt_id: string,
    turns: number = MAX_TURNS,
    isInvalidStreamRetry: boolean = false,
+    displayContent?: PartListUnion,
  ): AsyncGenerator<ServerGeminiStreamEvent, Turn> {
    if (!isInvalidStreamRetry) {
      this.config.resetTurn();
@@ -809,6 +818,7 @@ export class GeminiClient {
        prompt_id,
        boundedTurns,
        isInvalidStreamRetry,
+        displayContent,
      );

      // Fire AfterAgent hook if we have a turn and no pending tools
@@ -860,6 +870,8 @@ export class GeminiClient {
            signal,
            prompt_id,
            boundedTurns - 1,
+            false,
+            displayContent,
          );
        }
      }
--- a/packages/core/src/core/geminiChat.ts
+++ b/packages/core/src/core/geminiChat.ts
@@ -268,6 +268,7 @@ export class GeminiChat {
   * @param message - The list of messages to send.
   * @param prompt_id - The ID of the prompt.
   * @param signal - An abort signal for this message.
+   * @param displayContent - An optional user-friendly version of the message to record.
   * @return The model's response.
   *
   * @example
@@ -286,6 +287,7 @@ export class GeminiChat {
    message: PartListUnion,
    prompt_id: string,
    signal: AbortSignal,
+    displayContent?: PartListUnion,
  ): Promise<AsyncGenerator<StreamEvent>> {
    await this.sendPromise;

@@ -302,12 +304,25 @@ export class GeminiChat {
    // Record user input - capture complete message with all parts (text, files, images, etc.)
    // but skip recording function responses (tool call results) as they should be stored in tool call records
    if (!isFunctionResponse(userContent)) {
-      const userMessage = Array.isArray(message) ? message : [message];
-      const userMessageContent = partListUnionToString(toParts(userMessage));
+      const userMessageParts = userContent.parts || [];
+      const userMessageContent = partListUnionToString(userMessageParts);
+
+      let finalDisplayContent: Part[] | undefined = undefined;
+      if (displayContent !== undefined) {
+        const displayParts = toParts(
+          Array.isArray(displayContent) ? displayContent : [displayContent],
+        );
+        const displayContentString = partListUnionToString(displayParts);
+        if (displayContentString !== userMessageContent) {
+          finalDisplayContent = displayParts;
+        }
+      }
+
      this.chatRecordingService.recordMessage({
        model,
        type: 'user',
-        content: userMessageContent,
+        content: userMessageParts,
+        displayContent: finalDisplayContent,
      });
    }

--- a/packages/core/src/core/turn.test.ts
+++ b/packages/core/src/core/turn.test.ts
@@ -102,6 +102,7 @@ describe('Turn', () => {
        reqParts,
        'prompt-id-1',
        expect.any(AbortSignal),
+        undefined,
      );

      expect(events).toEqual([
--- a/packages/core/src/core/turn.ts
+++ b/packages/core/src/core/turn.ts
@@ -248,6 +248,7 @@ export class Turn {
    modelConfigKey: ModelConfigKey,
    req: PartListUnion,
    signal: AbortSignal,
+    displayContent?: PartListUnion,
  ): AsyncGenerator<ServerGeminiStreamEvent> {
    try {
      // Note: This assumes `sendMessageStream` yields events like
@@ -257,6 +258,7 @@ export class Turn {
        req,
        this.prompt_id,
        signal,
+        displayContent,
      );

      for await (const streamEvent of responseStream) {
--- a/packages/core/src/services/chatRecordingService.test.ts
+++ b/packages/core/src/services/chatRecordingService.test.ts
@@ -130,6 +130,7 @@ describe('ChatRecordingService', () => {
      chatRecordingService.recordMessage({
        type: 'user',
        content: 'Hello',
+        displayContent: 'User Hello',
        model: 'gemini-pro',
      });
      expect(mkdirSyncSpy).toHaveBeenCalled();
@@ -139,6 +140,7 @@ describe('ChatRecordingService', () => {
      ) as ConversationRecord;
      expect(conversation.messages).toHaveLength(1);
      expect(conversation.messages[0].content).toBe('Hello');
+      expect(conversation.messages[0].displayContent).toBe('User Hello');
      expect(conversation.messages[0].type).toBe('user');
    });

--- a/packages/core/src/services/chatRecordingService.ts
+++ b/packages/core/src/services/chatRecordingService.ts
@@ -47,6 +47,7 @@ export interface BaseMessageRecord {
  id: string;
  timestamp: string;
  content: PartListUnion;
+  displayContent?: PartListUnion;
 }

 /**
@@ -207,12 +208,14 @@ export class ChatRecordingService {
  private newMessage(
    type: ConversationRecordExtra['type'],
    content: PartListUnion,
+    displayContent?: PartListUnion,
  ): MessageRecord {
    return {
      id: randomUUID(),
      timestamp: new Date().toISOString(),
      type,
      content,
+      displayContent,
    };
  }

@@ -223,12 +226,17 @@ export class ChatRecordingService {
    model: string | undefined;
    type: ConversationRecordExtra['type'];
    content: PartListUnion;
+    displayContent?: PartListUnion;
  }): void {
    if (!this.conversationFile) return;

    try {
      this.updateConversation((conversation) => {
-        const msg = this.newMessage(message.type, message.content);
+        const msg = this.newMessage(
+          message.type,
+          message.content,
+          message.displayContent,
+        );
        if (msg.type === 'gemini') {
          // If it's a new Gemini message then incorporate any queued thoughts.
          conversation.messages.push({