From 57012ae5b33bcccee2e90e345447df5bbf653a3d Mon Sep 17 00:00:00 2001
From: Adib234 <30782825+Adib234@users.noreply.github.com>
Date: Wed, 7 Jan 2026 12:10:22 -0500
Subject: [PATCH] Core data structure updates for Rewind functionality (#15714)
---
.../components/messages/ToolMessage.test.tsx | 1 +
.../useToolScheduler.test.ts.snap | 1 +
packages/core/src/core/client.test.ts | 33 ++++++++++++
packages/core/src/core/client.ts | 14 ++++-
packages/core/src/core/coreToolScheduler.ts | 1 +
packages/core/src/core/geminiChat.ts | 5 +-
.../src/services/chatRecordingService.test.ts | 53 +++++++++++++++++++
.../core/src/services/chatRecordingService.ts | 35 ++++++++++--
packages/core/src/telemetry/loggers.test.ts | 1 +
packages/core/src/tools/edit.ts | 2 +
packages/core/src/tools/tools.ts | 2 +
packages/core/src/tools/write-file.ts | 2 +
12 files changed, 145 insertions(+), 5 deletions(-)
diff --git a/packages/cli/src/ui/components/messages/ToolMessage.test.tsx b/packages/cli/src/ui/components/messages/ToolMessage.test.tsx
index 983bca8669..fb01c4d9bc 100644
--- a/packages/cli/src/ui/components/messages/ToolMessage.test.tsx
+++ b/packages/cli/src/ui/components/messages/ToolMessage.test.tsx
@@ -182,6 +182,7 @@ describe('', () => {
fileName: 'file.txt',
originalContent: 'old',
newContent: 'new',
+ filePath: 'file.txt',
};
const { lastFrame } = renderWithContext(
,
diff --git a/packages/cli/src/ui/hooks/__snapshots__/useToolScheduler.test.ts.snap b/packages/cli/src/ui/hooks/__snapshots__/useToolScheduler.test.ts.snap
index 871be0e764..24ff4e1356 100644
--- a/packages/cli/src/ui/hooks/__snapshots__/useToolScheduler.test.ts.snap
+++ b/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,
},
diff --git a/packages/core/src/core/client.test.ts b/packages/core/src/core/client.test.ts
index 6045088c04..16f78d40d8 100644
--- a/packages/core/src/core/client.test.ts
+++ b/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', () => {
diff --git a/packages/core/src/core/client.ts b/packages/core/src/core/client.ts
index 48da7e43e7..bf70aa2200 100644
--- a/packages/core/src/core/client.ts
+++ b/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;
}
diff --git a/packages/core/src/core/coreToolScheduler.ts b/packages/core/src/core/coreToolScheduler.ts
index 69a2e03475..1120074248 100644
--- a/packages/core/src/core/coreToolScheduler.ts
+++ b/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,
};
}
}
diff --git a/packages/core/src/core/geminiChat.ts b/packages/core/src/core/geminiChat.ts
index 3dc91e1b6c..3bc928c6fb 100644
--- a/packages/core/src/core/geminiChat.ts
+++ b/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,
diff --git a/packages/core/src/services/chatRecordingService.test.ts b/packages/core/src/services/chatRecordingService.test.ts
index dcd77c986f..6fb49fbd5f 100644
--- a/packages/core/src/services/chatRecordingService.test.ts
+++ b/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();
+ });
+ });
});
diff --git a/packages/core/src/services/chatRecordingService.ts b/packages/core/src/services/chatRecordingService.ts
index 0f4dab0f49..b308cce789 100644
--- a/packages/core/src/services/chatRecordingService.ts
+++ b/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;
+ }
}
diff --git a/packages/core/src/telemetry/loggers.test.ts b/packages/core/src/telemetry/loggers.test.ts
index 3dabc4a89d..c0023a1680 100644
--- a/packages/core/src/telemetry/loggers.test.ts
+++ b/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: {
diff --git a/packages/core/src/tools/edit.ts b/packages/core/src/tools/edit.ts
index 3f71bdaad0..85c86ff804 100644
--- a/packages/core/src/tools/edit.ts
+++ b/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,
};
}
diff --git a/packages/core/src/tools/tools.ts b/packages/core/src/tools/tools.ts
index 1b6f6f92ee..d3efd56ec1 100644
--- a/packages/core/src/tools/tools.ts
+++ b/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 {
diff --git a/packages/core/src/tools/write-file.ts b/packages/core/src/tools/write-file.ts
index 339a60b4b6..3dbb696acc 100644
--- a/packages/core/src/tools/write-file.ts
+++ b/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 {