diff --git a/packages/cli/src/nonInteractiveCli.test.ts b/packages/cli/src/nonInteractiveCli.test.ts index 3217645211..d0e21b6b6d 100644 --- 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'); diff --git a/packages/cli/src/nonInteractiveCli.ts b/packages/cli/src/nonInteractiveCli.ts index 9535fbded2..a2ca92a4e8 100644 --- 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 = ''; diff --git a/packages/cli/src/ui/components/RewindViewer.test.tsx b/packages/cli/src/ui/components/RewindViewer.test.tsx index cdb0650408..5ad1f8b5e4 100644 --- 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()); + }); }); }); diff --git a/packages/cli/src/ui/components/RewindViewer.tsx b/packages/cli/src/ui/components/RewindViewer.tsx index 2ab417888a..7a6143a6eb 100644 --- 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 = ({ conversation, onExit, @@ -162,10 +170,7 @@ export const RewindViewer: React.FC = ({ (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 = ({ isSelected ? theme.status.success : theme.text.primary } > - {partToString(userPrompt.content)} + {partToString( + userPrompt.displayContent || userPrompt.content, + )} Cancel rewind and stay here @@ -235,10 +242,7 @@ export const RewindViewer: React.FC = ({ 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 ( diff --git a/packages/cli/src/ui/components/__snapshots__/RewindViewer.test.tsx.snap b/packages/cli/src/ui/components/__snapshots__/RewindViewer.test.tsx.snap index d5acee4a7b..9ae46a1e05 100644 --- 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`] = ` "╭──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╮ │ │ diff --git a/packages/cli/src/ui/hooks/useGeminiStream.test.tsx b/packages/cli/src/ui/hooks/useGeminiStream.test.tsx index b4971323cc..1c4434a34a 100644 --- 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', ); }); }); diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts index 100d069323..eca933d982 100644 --- 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, diff --git a/packages/cli/src/ui/hooks/useSessionBrowser.test.ts b/packages/cli/src/ui/hooks/useSessionBrowser.test.ts index a376d525c9..9aaf13c8ef 100644 --- 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, diff --git a/packages/cli/src/ui/hooks/useSessionBrowser.ts b/packages/cli/src/ui/hooks/useSessionBrowser.ts index 3d9619d738..c214011c8b 100644 --- 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 diff --git a/packages/core/src/core/client.test.ts b/packages/core/src/core/client.test.ts index ba30a63327..5d1edab256 100644 --- 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, ); }); diff --git a/packages/core/src/core/client.ts b/packages/core/src/core/client.ts index d80b8b4002..d6c3bb8520 100644 --- 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 { // 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 { 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, ); } } diff --git a/packages/core/src/core/geminiChat.ts b/packages/core/src/core/geminiChat.ts index bd7182fd03..a9cf192418 100644 --- 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> { 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, }); } diff --git a/packages/core/src/core/turn.test.ts b/packages/core/src/core/turn.test.ts index 43146e31ec..438ccdb55a 100644 --- 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([ diff --git a/packages/core/src/core/turn.ts b/packages/core/src/core/turn.ts index 8e6974704d..aa46c5d080 100644 --- 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 { 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) { diff --git a/packages/core/src/services/chatRecordingService.test.ts b/packages/core/src/services/chatRecordingService.test.ts index cba0a2e977..6dcfa79a77 100644 --- 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'); }); diff --git a/packages/core/src/services/chatRecordingService.ts b/packages/core/src/services/chatRecordingService.ts index a0a8034ce8..e570923d54 100644 --- 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({