Core data structure updates for Rewind functionality (#15714)

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;
+  }
 }