Core data structure updates for Rewind functionality (#15714)

packages/cli/src/ui/components/messages/ToolMessage.test.tsx

@@ -182,6 +182,7 @@ describe('<ToolMessage />', () => {
       fileName: 'file.txt',
       originalContent: 'old',
       newContent: 'new',
+      filePath: 'file.txt',
     };
     const { lastFrame } = renderWithContext(
       <ToolMessage {...baseProps} resultDisplay={diffResult} />,

packages/cli/src/ui/hooks/__snapshots__/useToolScheduler.test.ts.snap

@@ -64,6 +64,7 @@ exports[`useReactToolScheduler > should handle tool requiring confirmation - can
   "resultDisplay": {
     "fileDiff": "Mock tool requires confirmation",
     "fileName": "mockToolRequiresConfirmation.ts",
+    "filePath": undefined,
     "newContent": undefined,
     "originalContent": undefined,
   },

packages/core/src/core/client.test.ts

@@ -38,6 +38,7 @@ import { ideContextStore } from '../ide/ideContext.js';
 import type { ModelRouterService } from '../routing/modelRouterService.js';
 import { uiTelemetryService } from '../telemetry/uiTelemetry.js';
 import { ChatCompressionService } from '../services/chatCompressionService.js';
+import type { ChatRecordingService } from '../services/chatRecordingService.js';
 import { createAvailabilityServiceMock } from '../availability/testUtils.js';
 import type { ModelAvailabilityService } from '../availability/modelAvailabilityService.js';
 import type {
@@ -397,6 +398,10 @@ describe('Gemini Client (client.ts)', () => {
         getHistory: vi.fn((_curated?: boolean) => chatHistory),
         setHistory: vi.fn(),
         getLastPromptTokenCount: vi.fn().mockReturnValue(originalTokenCount),
+        getChatRecordingService: vi.fn().mockReturnValue({
+          getConversation: vi.fn().mockReturnValue(null),
+          getConversationFilePath: vi.fn().mockReturnValue(null),
+        }),
       };
       client['chat'] = mockOriginalChat as GeminiChat;
 
@@ -617,6 +622,34 @@ describe('Gemini Client (client.ts)', () => {
         newTokenCount: 50,
       });
     });
+
+    it('should resume the session file when compression succeeds', async () => {
+      const { client, mockOriginalChat } = setup({
+        compressionStatus: CompressionStatus.COMPRESSED,
+      });
+
+      const mockConversation = { some: 'conversation' };
+      const mockFilePath = '/tmp/session.json';
+
+      // Override the mock to return values
+      const mockRecordingService = {
+        getConversation: vi.fn().mockReturnValue(mockConversation),
+        getConversationFilePath: vi.fn().mockReturnValue(mockFilePath),
+      };
+      vi.mocked(mockOriginalChat.getChatRecordingService!).mockReturnValue(
+        mockRecordingService as unknown as ChatRecordingService,
+      );
+
+      await client.tryCompressChat('prompt-id', false);
+
+      expect(client['startChat']).toHaveBeenCalledWith(
+        expect.anything(), // newHistory
+        {
+          conversation: mockConversation,
+          filePath: mockFilePath,
+        },
+      );
+    });
   });
 
   describe('sendMessageStream', () => {

packages/core/src/core/client.ts

@@ -972,7 +972,19 @@ export class GeminiClient {
         this.hasFailedCompressionAttempt || !force;
     } else if (info.compressionStatus === CompressionStatus.COMPRESSED) {
       if (newHistory) {
-        this.chat = await this.startChat(newHistory);
+        // capture current session data before resetting
+        const currentRecordingService =
+          this.getChat().getChatRecordingService();
+        const conversation = currentRecordingService.getConversation();
+        const filePath = currentRecordingService.getConversationFilePath();
+
+        let resumedData: ResumedSessionData | undefined;
+
+        if (conversation && filePath) {
+          resumedData = { conversation, filePath };
+        }
+
+        this.chat = await this.startChat(newHistory, resumedData);
         this.updateTelemetryTokenCount();
         this.forceFullIdeContext = true;
       }

packages/core/src/core/coreToolScheduler.ts

@@ -279,6 +279,7 @@ export class CoreToolScheduler {
                 originalContent:
                   waitingCall.confirmationDetails.originalContent,
                 newContent: waitingCall.confirmationDetails.newContent,
+                filePath: waitingCall.confirmationDetails.filePath,
               };
             }
           }

packages/core/src/core/geminiChat.ts

@@ -838,7 +838,10 @@ export class GeminiChat {
     const toolCallRecords = toolCalls.map((call) => {
       const resultDisplayRaw = call.response?.resultDisplay;
       const resultDisplay =
-        typeof resultDisplayRaw === 'string' ? resultDisplayRaw : undefined;
+        typeof resultDisplayRaw === 'string' ||
+        (typeof resultDisplayRaw === 'object' && resultDisplayRaw !== null)
+          ? resultDisplayRaw
+          : undefined;
 
       return {
         id: call.request.callId,

packages/core/src/services/chatRecordingService.test.ts

@@ -401,4 +401,57 @@ describe('ChatRecordingService', () => {
       );
     });
   });
+
+  describe('rewindTo', () => {
+    it('should rewind the conversation to a specific message ID', () => {
+      chatRecordingService.initialize();
+      const initialConversation = {
+        sessionId: 'test-session-id',
+        projectHash: 'test-project-hash',
+        messages: [
+          { id: '1', type: 'user', content: 'msg1' },
+          { id: '2', type: 'gemini', content: 'msg2' },
+          { id: '3', type: 'user', content: 'msg3' },
+        ],
+      };
+      vi.spyOn(fs, 'readFileSync').mockReturnValue(
+        JSON.stringify(initialConversation),
+      );
+      const writeFileSyncSpy = vi
+        .spyOn(fs, 'writeFileSync')
+        .mockImplementation(() => undefined);
+
+      const result = chatRecordingService.rewindTo('2');
+
+      if (!result) throw new Error('Result should not be null');
+      expect(result.messages).toHaveLength(1);
+      expect(result.messages[0].id).toBe('1');
+      expect(writeFileSyncSpy).toHaveBeenCalled();
+      const savedConversation = JSON.parse(
+        writeFileSyncSpy.mock.calls[0][1] as string,
+      ) as ConversationRecord;
+      expect(savedConversation.messages).toHaveLength(1);
+    });
+
+    it('should return the original conversation if the message ID is not found', () => {
+      chatRecordingService.initialize();
+      const initialConversation = {
+        sessionId: 'test-session-id',
+        projectHash: 'test-project-hash',
+        messages: [{ id: '1', type: 'user', content: 'msg1' }],
+      };
+      vi.spyOn(fs, 'readFileSync').mockReturnValue(
+        JSON.stringify(initialConversation),
+      );
+      const writeFileSyncSpy = vi
+        .spyOn(fs, 'writeFileSync')
+        .mockImplementation(() => undefined);
+
+      const result = chatRecordingService.rewindTo('non-existent');
+
+      if (!result) throw new Error('Result should not be null');
+      expect(result.messages).toHaveLength(1);
+      expect(writeFileSyncSpy).not.toHaveBeenCalled();
+    });
+  });
 });

packages/core/src/services/chatRecordingService.ts

@@ -16,6 +16,7 @@ import type {
   GenerateContentResponseUsageMetadata,
 } from '@google/genai';
 import { debugLogger } from '../utils/debugLogger.js';
+import type { ToolResultDisplay } from '../tools/tools.js';
 
 export const SESSION_FILE_PREFIX = 'session-';
 
@@ -53,7 +54,7 @@ export interface ToolCallRecord {
   // UI-specific fields for display purposes
   displayName?: string;
   description?: string;
-  resultDisplay?: string;
+  resultDisplay?: ToolResultDisplay;
   renderOutputAsMarkdown?: boolean;
 }
 
@@ -407,11 +408,14 @@ export class ChatRecordingService {
   /**
    * Saves the conversation record; overwrites the file.
    */
-  private writeConversation(conversation: ConversationRecord): void {
+  private writeConversation(
+    conversation: ConversationRecord,
+    { allowEmpty = false }: { allowEmpty?: boolean } = {},
+  ): void {
     try {
       if (!this.conversationFile) return;
       // Don't write the file yet until there's at least one message.
-      if (conversation.messages.length === 0) return;
+      if (conversation.messages.length === 0 && !allowEmpty) return;
 
       // Only write the file if this change would change the file.
       if (this.cachedLastConvData !== JSON.stringify(conversation, null, 2)) {
@@ -492,4 +496,29 @@ export class ChatRecordingService {
       throw error;
     }
   }
+
+  /**
+   * Rewinds the conversation to the state just before the specified message ID.
+   * All messages from (and including) the specified ID onwards are removed.
+   */
+  rewindTo(messageId: string): ConversationRecord | null {
+    if (!this.conversationFile) {
+      return null;
+    }
+    const conversation = this.readConversation();
+    const messageIndex = conversation.messages.findIndex(
+      (m) => m.id === messageId,
+    );
+
+    if (messageIndex === -1) {
+      debugLogger.error(
+        'Message to rewind to not found in conversation history',
+      );
+      return conversation;
+    }
+
+    conversation.messages = conversation.messages.slice(0, messageIndex);
+    this.writeConversation(conversation, { allowEmpty: true });
+    return conversation;
+  }
 }

packages/core/src/telemetry/loggers.test.ts

@@ -1053,6 +1053,7 @@ describe('loggers', () => {
           resultDisplay: {
             fileDiff: 'diff',
             fileName: 'file.txt',
+            filePath: 'file.txt',
             originalContent: 'old content',
             newContent: 'new content',
             diffStat: {

packages/core/src/tools/edit.ts

@@ -818,9 +818,11 @@ class EditToolInvocation
         displayResult = {
           fileDiff,
           fileName,
+          filePath: this.params.file_path,
           originalContent: editData.currentContent,
           newContent: editData.newContent,
           diffStat,
+          isNewFile: editData.isNewFile,
         };
       }
 

packages/core/src/tools/tools.ts

@@ -647,9 +647,11 @@ export interface Todo {
 export interface FileDiff {
   fileDiff: string;
   fileName: string;
+  filePath: string;
   originalContent: string | null;
   newContent: string;
   diffStat?: DiffStat;
+  isNewFile?: boolean;
 }
 
 export interface DiffStat {

packages/core/src/tools/write-file.ts

@@ -346,9 +346,11 @@ class WriteFileToolInvocation extends BaseToolInvocation<
       const displayResult: FileDiff = {
         fileDiff,
         fileName,
+        filePath: this.resolvedPath,
         originalContent: correctedContentResult.originalContent,
         newContent: correctedContentResult.correctedContent,
         diffStat,
+        isNewFile,
       };
 
       return {