From 6f4382901e473b7a530f22157e52dacfb45147c6 Mon Sep 17 00:00:00 2001 From: Abhijit Balaji Date: Thu, 5 Mar 2026 17:13:34 -0800 Subject: [PATCH] fix(core): ensure chat compression summary persists correctly without data loss - Modified ChatRecordingService.initialize to accept an explicit overwriteHistory flag. - When overwriteHistory is true (e.g., after chat compression), it overwrites disk messages; otherwise, it preserves existing historical metadata (IDs/timestamps). - Updated GeminiChat and GeminiClient.startChat to propagate this flag, ensuring it is only true during the compression flow. - Refactored apiContentToMessageRecords for full type safety by removing unsafe type assertions and adding part fallbacks, as suggested in PR review. - Updated unit tests in chatRecordingService.test.ts and client.test.ts to verify the new behavior and fix regression. - Verified all workspace and integration tests pass via preflight. Fixes #21335 --- packages/core/src/core/client.test.ts | 1 + packages/core/src/core/client.ts | 5 +- packages/core/src/core/geminiChat.ts | 8 +++- .../src/services/chatRecordingService.test.ts | 46 ++++++++++++++++++- .../core/src/services/chatRecordingService.ts | 12 +++-- 5 files changed, 65 insertions(+), 7 deletions(-) diff --git a/packages/core/src/core/client.test.ts b/packages/core/src/core/client.test.ts index 059b72437f..4a6e7d2825 100644 --- a/packages/core/src/core/client.test.ts +++ b/packages/core/src/core/client.test.ts @@ -661,6 +661,7 @@ describe('Gemini Client (client.ts)', () => { conversation: mockConversation, filePath: mockFilePath, }, + true, // overwriteHistory ); }); }); diff --git a/packages/core/src/core/client.ts b/packages/core/src/core/client.ts index 14e2f42bc3..5a3f79fc01 100644 --- a/packages/core/src/core/client.ts +++ b/packages/core/src/core/client.ts @@ -335,6 +335,7 @@ export class GeminiClient { async startChat( extraHistory?: Content[], resumedSessionData?: ResumedSessionData, + overwriteHistory: boolean = false, ): Promise { this.forceFullIdeContext = true; this.hasFailedCompressionAttempt = false; @@ -362,6 +363,8 @@ export class GeminiClient { toolRegistry.getFunctionDeclarations(modelId); return [{ functionDeclarations: toolDeclarations }]; }, + 'main', + overwriteHistory, ); } catch (error) { await reportError( @@ -1136,7 +1139,7 @@ export class GeminiClient { resumedData = { conversation, filePath }; } - this.chat = await this.startChat(newHistory, resumedData); + this.chat = await this.startChat(newHistory, resumedData, true); this.updateTelemetryTokenCount(); this.forceFullIdeContext = true; } diff --git a/packages/core/src/core/geminiChat.ts b/packages/core/src/core/geminiChat.ts index df522c6eb4..184a1785df 100644 --- a/packages/core/src/core/geminiChat.ts +++ b/packages/core/src/core/geminiChat.ts @@ -252,10 +252,16 @@ export class GeminiChat { resumedSessionData?: ResumedSessionData, private readonly onModelChanged?: (modelId: string) => Promise, kind: 'main' | 'subagent' = 'main', + overwriteHistory: boolean = false, ) { validateHistory(history); this.chatRecordingService = new ChatRecordingService(config); - this.chatRecordingService.initialize(resumedSessionData, kind, history); + this.chatRecordingService.initialize( + resumedSessionData, + kind, + history, + overwriteHistory, + ); this.lastPromptTokenCount = estimateTokenCountSync( this.history.flatMap((c) => c.parts || []), ); diff --git a/packages/core/src/services/chatRecordingService.test.ts b/packages/core/src/services/chatRecordingService.test.ts index bc2e9375a2..be69e7a71a 100644 --- a/packages/core/src/services/chatRecordingService.test.ts +++ b/packages/core/src/services/chatRecordingService.test.ts @@ -123,7 +123,7 @@ describe('ChatRecordingService', () => { expect(conversation.sessionId).toBe('old-session-id'); }); - it('should overwrite existing messages if initialHistory is provided', () => { + it('should overwrite existing messages if overwriteHistory is true', () => { const chatsDir = path.join(testTempDir, 'chats'); fs.mkdirSync(chatsDir, { recursive: true }); const sessionFile = path.join(chatsDir, 'session-overwrite.json'); @@ -155,6 +155,7 @@ describe('ChatRecordingService', () => { }, 'main', newHistory, + true, // overwriteHistory = true ); const conversation = JSON.parse( @@ -166,6 +167,49 @@ describe('ChatRecordingService', () => { ]); expect(conversation.messages[0].type).toBe('user'); }); + + it('should NOT overwrite existing messages if overwriteHistory is false', () => { + const chatsDir = path.join(testTempDir, 'chats'); + fs.mkdirSync(chatsDir, { recursive: true }); + const sessionFile = path.join(chatsDir, 'session-no-overwrite.json'); + const initialData = { + sessionId: 'test-session-id', + projectHash: 'test-project-hash', + messages: [ + { + id: 'msg-1', + type: 'user', + content: 'Old Message', + timestamp: new Date().toISOString(), + }, + ], + }; + fs.writeFileSync(sessionFile, JSON.stringify(initialData)); + + const newHistory: Content[] = [ + { + role: 'user', + parts: [{ text: 'Some New Context' }], + }, + ]; + + chatRecordingService.initialize( + { + filePath: sessionFile, + conversation: initialData as ConversationRecord, + }, + 'main', + newHistory, + false, // overwriteHistory = false + ); + + const conversation = JSON.parse( + fs.readFileSync(sessionFile, 'utf8'), + ) as ConversationRecord; + expect(conversation.messages).toHaveLength(1); + expect(conversation.messages[0].content).toBe('Old Message'); + expect(conversation.messages[0].id).toBe('msg-1'); + }); }); describe('recordMessage', () => { diff --git a/packages/core/src/services/chatRecordingService.ts b/packages/core/src/services/chatRecordingService.ts index ed589113a7..6381bebff0 100644 --- a/packages/core/src/services/chatRecordingService.ts +++ b/packages/core/src/services/chatRecordingService.ts @@ -147,14 +147,16 @@ export class ChatRecordingService { * * @param resumedSessionData Data from a previous session to resume from. * @param kind The kind of conversation (main or subagent). - * @param initialHistory The starting history for this session. If provided when resuming an - * existing session (e.g., after chat compression), this will overwrite the messages currently - * stored on disk to ensure the file reflects the new session state. + * @param initialHistory The starting history for this session. + * @param overwriteHistory If true and resuming an existing session, the messages on disk will + * be overwritten with initialHistory. Use with caution as this destroys original message + * IDs and timestamps. */ initialize( resumedSessionData?: ResumedSessionData, kind?: 'main' | 'subagent', initialHistory?: Content[], + overwriteHistory: boolean = false, ): void { try { this.kind = kind; @@ -167,7 +169,7 @@ export class ChatRecordingService { // Update the session ID in the existing file this.updateConversation((conversation) => { conversation.sessionId = this.sessionId; - if (initialHistory) { + if (overwriteHistory && initialHistory) { conversation.messages = this.apiContentToMessageRecords(initialHistory); } @@ -225,6 +227,8 @@ export class ChatRecordingService { /** * Converts API Content array to storage-compatible MessageRecord array. + * WARNING: Generates new IDs and timestamps. Use only for brand new content + * or when explicitly overwriting history (e.g. after compression). */ private apiContentToMessageRecords(history: Content[]): MessageRecord[] { return history.map((content): MessageRecord => {