feat(core): implement persistence and resumption for masked tool outputs (#18451)

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 {

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 =
+        '<tool_output_masked>short preview</tool_output_masked>';
+      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 = '<masked>';
+      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);
+    });
+  });
 });

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<string, Part[]>();
+        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;
+    }
+  }
 }