From 764a75fb3d15e7bbcc0bb28bb601a1dbfbd157b5 Mon Sep 17 00:00:00 2001 From: Abhi Date: Fri, 6 Feb 2026 12:56:34 -0500 Subject: [PATCH 1/2] feat(core): implement tracking and persistence for masked tool outputs This PR enables the synchronization of masked tool outputs to the session history, ensuring that automatic and manual masking updates are persisted to disk and correctly resumed in future sessions. It also fixes a bug where masked outputs were resumed as plain text parts, breaking the tool-use trajectory and stripping multi-modal data. --- packages/core/src/core/geminiChat.ts | 1 + .../src/services/chatRecordingService.test.ts | 143 ++++++++++++++++++ .../core/src/services/chatRecordingService.ts | 57 +++++++ 3 files changed, 201 insertions(+) diff --git a/packages/core/src/core/geminiChat.ts b/packages/core/src/core/geminiChat.ts index 69c494a4e0c..e053dd4c676 100644 --- a/packages/core/src/core/geminiChat.ts +++ b/packages/core/src/core/geminiChat.ts @@ -705,6 +705,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 e8b879e10ce..748b3a6adf0 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,146 @@ 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'); + }); + }); }); diff --git a/packages/core/src/services/chatRecordingService.ts b/packages/core/src/services/chatRecordingService.ts index 6a57e2801bf..435619e31f2 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,59 @@ 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) { + let currentCallId: string | undefined; + for (const part of content.parts) { + if (part.functionResponse && part.functionResponse.id) { + currentCallId = part.functionResponse.id; + if (!partsMap.has(currentCallId)) { + partsMap.set(currentCallId, []); + } + partsMap.get(currentCallId)!.push(part); + } else if (currentCallId) { + // This handles sibling parts (e.g. inlineData) for the same tool call in Gemini 2.x + 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; + } + } } From 223fd28eef19c00c722388e7ef5bf8d931e524a7 Mon Sep 17 00:00:00 2001 From: Abhi Date: Fri, 6 Feb 2026 16:09:54 -0500 Subject: [PATCH 2/2] fix brittle ordering pattern --- .../src/services/chatRecordingService.test.ts | 53 +++++++++++++++++++ .../core/src/services/chatRecordingService.ts | 25 +++++---- 2 files changed, 69 insertions(+), 9 deletions(-) diff --git a/packages/core/src/services/chatRecordingService.test.ts b/packages/core/src/services/chatRecordingService.test.ts index 748b3a6adf0..28d458c14b6 100644 --- a/packages/core/src/services/chatRecordingService.test.ts +++ b/packages/core/src/services/chatRecordingService.test.ts @@ -690,5 +690,58 @@ describe('ChatRecordingService', () => { 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 435619e31f2..ebe66edf01c 100644 --- a/packages/core/src/services/chatRecordingService.ts +++ b/packages/core/src/services/chatRecordingService.ts @@ -612,18 +612,25 @@ export class ChatRecordingService { const partsMap = new Map(); for (const content of history) { if (content.role === 'user' && content.parts) { - let currentCallId: string | undefined; + // 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 && part.functionResponse.id) { + if (part.functionResponse?.id) { currentCallId = part.functionResponse.id; - if (!partsMap.has(currentCallId)) { - partsMap.set(currentCallId, []); - } - partsMap.get(currentCallId)!.push(part); - } else if (currentCallId) { - // This handles sibling parts (e.g. inlineData) for the same tool call in Gemini 2.x - partsMap.get(currentCallId)!.push(part); } + + if (!partsMap.has(currentCallId)) { + partsMap.set(currentCallId, []); + } + partsMap.get(currentCallId)!.push(part); } } }