From 63f7e307905c14e7af71332c8e44b20afe32cc05 Mon Sep 17 00:00:00 2001 From: Abhi <43648792+abhipatel12@users.noreply.github.com> Date: Fri, 6 Feb 2026 16:22:22 -0500 Subject: [PATCH] feat(core): implement persistence and resumption for masked tool outputs (#18451) --- packages/core/src/core/geminiChat.ts | 1 + .../src/services/chatRecordingService.test.ts | 196 ++++++++++++++++++ .../core/src/services/chatRecordingService.ts | 64 ++++++ 3 files changed, 261 insertions(+) diff --git a/packages/core/src/core/geminiChat.ts b/packages/core/src/core/geminiChat.ts index c45642c7be..df98e3ebd7 100644 --- a/packages/core/src/core/geminiChat.ts +++ b/packages/core/src/core/geminiChat.ts @@ -699,6 +699,7 @@ export class GeminiChat { this.lastPromptTokenCount = estimateTokenCountSync( this.history.flatMap((c) => c.parts || []), ); + this.chatRecordingService.updateMessagesFromHistory(history); } stripThoughtsFromHistory(): void { diff --git a/packages/core/src/services/chatRecordingService.test.ts b/packages/core/src/services/chatRecordingService.test.ts index e8b879e10c..28d458c14b 100644 --- a/packages/core/src/services/chatRecordingService.test.ts +++ b/packages/core/src/services/chatRecordingService.test.ts @@ -13,6 +13,7 @@ import type { ToolCallRecord, MessageRecord, } from './chatRecordingService.js'; +import type { Content, Part } from '@google/genai'; import { ChatRecordingService } from './chatRecordingService.js'; import type { Config } from '../config/config.js'; import { getProjectHash } from '../utils/paths.js'; @@ -548,4 +549,199 @@ describe('ChatRecordingService', () => { writeFileSyncSpy.mockRestore(); }); }); + + describe('updateMessagesFromHistory', () => { + beforeEach(() => { + chatRecordingService.initialize(); + }); + + it('should update tool results from API history (masking sync)', () => { + // 1. Record an initial message and tool call + chatRecordingService.recordMessage({ + type: 'gemini', + content: 'I will list the files.', + model: 'gemini-pro', + }); + + const callId = 'tool-call-123'; + const originalResult = [{ text: 'a'.repeat(1000) }]; + chatRecordingService.recordToolCalls('gemini-pro', [ + { + id: callId, + name: 'list_files', + args: { path: '.' }, + result: originalResult, + status: 'success', + timestamp: new Date().toISOString(), + }, + ]); + + // 2. Prepare mock history with masked content + const maskedSnippet = + 'short preview'; + const history: Content[] = [ + { + role: 'model', + parts: [ + { functionCall: { name: 'list_files', args: { path: '.' } } }, + ], + }, + { + role: 'user', + parts: [ + { + functionResponse: { + name: 'list_files', + id: callId, + response: { output: maskedSnippet }, + }, + }, + ], + }, + ]; + + // 3. Trigger sync + chatRecordingService.updateMessagesFromHistory(history); + + // 4. Verify disk content + const sessionFile = chatRecordingService.getConversationFilePath()!; + const conversation = JSON.parse( + fs.readFileSync(sessionFile, 'utf8'), + ) as ConversationRecord; + + const geminiMsg = conversation.messages[0]; + if (geminiMsg.type !== 'gemini') + throw new Error('Expected gemini message'); + expect(geminiMsg.toolCalls).toBeDefined(); + expect(geminiMsg.toolCalls![0].id).toBe(callId); + // The implementation stringifies the response object + const result = geminiMsg.toolCalls![0].result; + if (!Array.isArray(result)) throw new Error('Expected array result'); + const firstPart = result[0] as Part; + expect(firstPart.functionResponse).toBeDefined(); + expect(firstPart.functionResponse!.id).toBe(callId); + expect(firstPart.functionResponse!.response).toEqual({ + output: maskedSnippet, + }); + }); + it('should preserve multi-modal sibling parts during sync', () => { + chatRecordingService.initialize(); + const callId = 'multi-modal-call'; + const originalResult: Part[] = [ + { + functionResponse: { + id: callId, + name: 'read_file', + response: { content: '...' }, + }, + }, + { inlineData: { mimeType: 'image/png', data: 'base64...' } }, + ]; + + chatRecordingService.recordMessage({ + type: 'gemini', + content: '', + model: 'gemini-pro', + }); + + chatRecordingService.recordToolCalls('gemini-pro', [ + { + id: callId, + name: 'read_file', + args: { path: 'image.png' }, + result: originalResult, + status: 'success', + timestamp: new Date().toISOString(), + }, + ]); + + const maskedSnippet = ''; + const history: Content[] = [ + { + role: 'user', + parts: [ + { + functionResponse: { + name: 'read_file', + id: callId, + response: { output: maskedSnippet }, + }, + }, + { inlineData: { mimeType: 'image/png', data: 'base64...' } }, + ], + }, + ]; + + chatRecordingService.updateMessagesFromHistory(history); + + const sessionFile = chatRecordingService.getConversationFilePath()!; + const conversation = JSON.parse( + fs.readFileSync(sessionFile, 'utf8'), + ) as ConversationRecord; + + const lastMsg = conversation.messages[0] as MessageRecord & { + type: 'gemini'; + }; + const result = lastMsg.toolCalls![0].result as Part[]; + expect(result).toHaveLength(2); + expect(result[0].functionResponse!.response).toEqual({ + output: maskedSnippet, + }); + expect(result[1].inlineData).toBeDefined(); + expect(result[1].inlineData!.mimeType).toBe('image/png'); + }); + + it('should handle parts appearing BEFORE the functionResponse in a content block', () => { + chatRecordingService.initialize(); + const callId = 'prefix-part-call'; + + chatRecordingService.recordMessage({ + type: 'gemini', + content: '', + model: 'gemini-pro', + }); + + chatRecordingService.recordToolCalls('gemini-pro', [ + { + id: callId, + name: 'read_file', + args: { path: 'test.txt' }, + result: [], + status: 'success', + timestamp: new Date().toISOString(), + }, + ]); + + const history: Content[] = [ + { + role: 'user', + parts: [ + { text: 'Prefix metadata or text' }, + { + functionResponse: { + name: 'read_file', + id: callId, + response: { output: 'file content' }, + }, + }, + ], + }, + ]; + + chatRecordingService.updateMessagesFromHistory(history); + + const sessionFile = chatRecordingService.getConversationFilePath()!; + const conversation = JSON.parse( + fs.readFileSync(sessionFile, 'utf8'), + ) as ConversationRecord; + + const lastMsg = conversation.messages[0] as MessageRecord & { + type: 'gemini'; + }; + const result = lastMsg.toolCalls![0].result as Part[]; + expect(result).toHaveLength(2); + expect(result[0].text).toBe('Prefix metadata or text'); + expect(result[1].functionResponse!.id).toBe(callId); + }); + }); }); diff --git a/packages/core/src/services/chatRecordingService.ts b/packages/core/src/services/chatRecordingService.ts index 6a57e2801b..ebe66edf01 100644 --- a/packages/core/src/services/chatRecordingService.ts +++ b/packages/core/src/services/chatRecordingService.ts @@ -13,6 +13,8 @@ import path from 'node:path'; import fs from 'node:fs'; import { randomUUID } from 'node:crypto'; import type { + Content, + Part, PartListUnion, GenerateContentResponseUsageMetadata, } from '@google/genai'; @@ -594,4 +596,66 @@ export class ChatRecordingService { this.writeConversation(conversation, { allowEmpty: true }); return conversation; } + + /** + * Updates the conversation history based on the provided API Content array. + * This is used to persist changes made to the history (like masking) back to disk. + */ + updateMessagesFromHistory(history: Content[]): void { + if (!this.conversationFile) return; + + try { + this.updateConversation((conversation) => { + // Create a map of tool results from the API history for quick lookup by call ID. + // We store the full list of parts associated with each tool call ID to preserve + // multi-modal data and proper trajectory structure. + const partsMap = new Map(); + for (const content of history) { + if (content.role === 'user' && content.parts) { + // Find all unique call IDs in this message + const callIds = content.parts + .map((p) => p.functionResponse?.id) + .filter((id): id is string => !!id); + + if (callIds.length === 0) continue; + + // Use the first ID as a seed to capture any "leading" non-ID parts + // in this specific content block. + let currentCallId = callIds[0]; + for (const part of content.parts) { + if (part.functionResponse?.id) { + currentCallId = part.functionResponse.id; + } + + if (!partsMap.has(currentCallId)) { + partsMap.set(currentCallId, []); + } + partsMap.get(currentCallId)!.push(part); + } + } + } + + // Update the conversation records tool results if they've changed. + for (const message of conversation.messages) { + if (message.type === 'gemini' && message.toolCalls) { + for (const toolCall of message.toolCalls) { + const newParts = partsMap.get(toolCall.id); + if (newParts !== undefined) { + // Store the results as proper Parts (including functionResponse) + // instead of stringifying them as text parts. This ensures the + // tool trajectory is correctly reconstructed upon session resumption. + toolCall.result = newParts; + } + } + } + } + }); + } catch (error) { + debugLogger.error( + 'Error updating conversation history from memory.', + error, + ); + throw error; + } + } }