fix(context): Fix snapshot recovery across sessions. (#26939)

packages/core/src/services/chatRecordingService.test.ts

@@ -47,9 +47,10 @@ import {
 } from './chatRecordingService.js';
 import type { WorkspaceContext } from '../utils/workspaceContext.js';
 import { CoreToolCallStatus } from '../scheduler/types.js';
-import type { Content, Part } from '@google/genai';
+import type { Part } from '@google/genai';
 import type { Config } from '../config/config.js';
 import { getProjectHash } from '../utils/paths.js';
+import type { HistoryTurn } from '../core/agentChatHistory.js';
 
 vi.mock('../utils/paths.js');
 vi.mock('node:crypto', async (importOriginal) => {
@@ -1065,7 +1066,7 @@ describe('ChatRecordingService', () => {
 
     it('should update tool results from API history (masking sync)', async () => {
       // 1. Record an initial message and tool call
-      chatRecordingService.recordMessage({
+      const modelMsgId = chatRecordingService.recordMessage({
         type: 'gemini',
         content: 'I will list the files.',
         model: 'gemini-pro',
@@ -1087,24 +1088,30 @@ describe('ChatRecordingService', () => {
       // 2. Prepare mock history with masked content
       const maskedSnippet =
         '<tool_output_masked>short preview</tool_output_masked>';
-      const history: Content[] = [
+      const history: HistoryTurn[] = [
         {
-          role: 'model',
-          parts: [
-            { functionCall: { name: 'list_files', args: { path: '.' } } },
-          ],
+          id: modelMsgId,
+          content: {
+            role: 'model',
+            parts: [
+              { functionCall: { name: 'list_files', args: { path: '.' } } },
+            ],
+          },
         },
         {
-          role: 'user',
-          parts: [
-            {
-              functionResponse: {
-                name: 'list_files',
-                id: callId,
-                response: { output: maskedSnippet },
+          id: 'user-id',
+          content: {
+            role: 'user',
+            parts: [
+              {
+                functionResponse: {
+                  name: 'list_files',
+                  id: callId,
+                  response: { output: maskedSnippet },
+                },
               },
-            },
-          ],
+            ],
+          },
         },
       ];
 
@@ -1132,8 +1139,15 @@ describe('ChatRecordingService', () => {
         output: maskedSnippet,
       });
     });
+
     it('should preserve multi-modal sibling parts during sync', async () => {
       await chatRecordingService.initialize();
+      const modelMsgId = chatRecordingService.recordMessage({
+        type: 'gemini',
+        content: '',
+        model: 'gemini-pro',
+      });
+
       const callId = 'multi-modal-call';
       const originalResult: Part[] = [
         {
@@ -1146,12 +1160,6 @@ describe('ChatRecordingService', () => {
         { inlineData: { mimeType: 'image/png', data: 'base64...' } },
       ];
 
-      chatRecordingService.recordMessage({
-        type: 'gemini',
-        content: '',
-        model: 'gemini-pro',
-      });
-
       chatRecordingService.recordToolCalls('gemini-pro', [
         {
           id: callId,
@@ -1164,19 +1172,26 @@ describe('ChatRecordingService', () => {
       ]);
 
       const maskedSnippet = '<masked>';
-      const history: Content[] = [
+      const history: HistoryTurn[] = [
         {
-          role: 'user',
-          parts: [
-            {
-              functionResponse: {
-                name: 'read_file',
-                id: callId,
-                response: { output: maskedSnippet },
+          id: modelMsgId,
+          content: { role: 'model', parts: [] },
+        },
+        {
+          id: 'user-id',
+          content: {
+            role: 'user',
+            parts: [
+              {
+                functionResponse: {
+                  name: 'read_file',
+                  id: callId,
+                  response: { output: maskedSnippet },
+                },
               },
-            },
-            { inlineData: { mimeType: 'image/png', data: 'base64...' } },
-          ],
+              { inlineData: { mimeType: 'image/png', data: 'base64...' } },
+            ],
+          },
         },
       ];
 
@@ -1201,14 +1216,14 @@ describe('ChatRecordingService', () => {
 
     it('should handle parts appearing BEFORE the functionResponse in a content block', async () => {
       await chatRecordingService.initialize();
-      const callId = 'prefix-part-call';
-
-      chatRecordingService.recordMessage({
+      const modelMsgId = chatRecordingService.recordMessage({
         type: 'gemini',
         content: '',
         model: 'gemini-pro',
       });
 
+      const callId = 'prefix-part-call';
+
       chatRecordingService.recordToolCalls('gemini-pro', [
         {
           id: callId,
@@ -1220,19 +1235,26 @@ describe('ChatRecordingService', () => {
         },
       ]);
 
-      const history: Content[] = [
+      const history: HistoryTurn[] = [
         {
-          role: 'user',
-          parts: [
-            { text: 'Prefix metadata or text' },
-            {
-              functionResponse: {
-                name: 'read_file',
-                id: callId,
-                response: { output: 'file content' },
+          id: modelMsgId,
+          content: { role: 'model', parts: [] },
+        },
+        {
+          id: 'user-id',
+          content: {
+            role: 'user',
+            parts: [
+              { text: 'Prefix metadata or text' },
+              {
+                functionResponse: {
+                  name: 'read_file',
+                  id: callId,
+                  response: { output: 'file content' },
+                },
               },
-            },
-          ],
+            ],
+          },
         },
       ];
 
@@ -1263,25 +1285,30 @@ describe('ChatRecordingService', () => {
       appendFileSyncSpy.mockClear();
 
       // History with a tool call ID that doesn't exist in the conversation
-      const history: Content[] = [
+      const history: HistoryTurn[] = [
         {
-          role: 'user',
-          parts: [
-            {
-              functionResponse: {
-                name: 'read_file',
-                id: 'nonexistent-call-id',
-                response: { output: 'some content' },
+          id: 'user-id',
+          content: {
+            role: 'user',
+            parts: [
+              {
+                functionResponse: {
+                  name: 'read_file',
+                  id: 'nonexistent-call-id',
+                  response: { output: 'some content' },
+                },
               },
-            },
-          ],
+            ],
+          },
         },
       ];
 
       chatRecordingService.updateMessagesFromHistory(history);
 
-      // No tool calls matched, so writeFileSync should NOT have been called
-      expect(appendFileSyncSpy).not.toHaveBeenCalled();
+      // In the new 'Strong Owner' architecture, updateMessagesFromHistory ensures that
+      // all turns in history (including new/synthetic ones) are recorded.
+      // Since 'user-id' was not in the original conversation, it is added.
+      expect(appendFileSyncSpy).toHaveBeenCalled();
     });
   });
 
@@ -1315,4 +1342,69 @@ describe('ChatRecordingService', () => {
       mkdirSyncSpy.mockRestore();
     });
   });
+
+  describe('recordSyntheticMessage and history sync', () => {
+    it('should correctly record synthetic messages with durable IDs', async () => {
+      await chatRecordingService.initialize(undefined, 'main');
+      const parts = [{ text: 'Synthetic Turn' }];
+
+      // Implicit ID generation
+      const id1 = chatRecordingService.recordSyntheticMessage('user', parts);
+      expect(id1).toBeDefined();
+      expect(id1).toMatch(/test-uuid-/);
+
+      // Explicit ID registration (e.g. from context processor)
+      const customId = 'stable-hash-123';
+      const id2 = chatRecordingService.recordSyntheticMessage(
+        'gemini',
+        parts,
+        customId,
+      );
+      expect(id2).toBe(customId);
+
+      const record = await loadConversationRecord(
+        chatRecordingService.getConversationFilePath()!,
+      );
+      expect(record!.messages).toHaveLength(2);
+      expect(record!.messages[0].id).toBe(id1);
+      expect(record!.messages[0].type).toBe('user');
+      expect(record!.messages[1].id).toBe(customId);
+      expect(record!.messages[1].type).toBe('gemini');
+    });
+
+    it('should synchronize history turns and maintain their durable identity', async () => {
+      await chatRecordingService.initialize(undefined, 'main');
+      const history: HistoryTurn[] = [
+        { id: 'h1', content: { role: 'user', parts: [{ text: 'msg1' }] } },
+        { id: 'h2', content: { role: 'model', parts: [{ text: 'msg2' }] } },
+      ];
+
+      chatRecordingService.updateMessagesFromHistory(history);
+
+      const record = await loadConversationRecord(
+        chatRecordingService.getConversationFilePath()!,
+      );
+      expect(record!.messages).toHaveLength(2);
+      expect(record!.messages[0].id).toBe('h1');
+      expect(record!.messages[1].id).toBe('h2');
+
+      // Update with a summary
+      const summaryId = 'summary-123';
+      const updatedHistory: HistoryTurn[] = [
+        {
+          id: summaryId,
+          content: { role: 'user', parts: [{ text: 'summary' }] },
+        },
+        ...history.slice(1),
+      ];
+
+      chatRecordingService.updateMessagesFromHistory(updatedHistory);
+      const record2 = await loadConversationRecord(
+        chatRecordingService.getConversationFilePath()!,
+      );
+      expect(record2!.messages).toHaveLength(2);
+      expect(record2!.messages[0].id).toBe(summaryId);
+      expect(record2!.messages[1].id).toBe('h2');
+    });
+  });
 });

packages/core/src/services/chatRecordingService.ts

@@ -17,13 +17,12 @@ import {
 import readline from 'node:readline';
 import { randomUUID } from 'node:crypto';
 import type {
-  Content,
-  Part,
   PartListUnion,
   GenerateContentResponseUsageMetadata,
 } from '@google/genai';
 import { debugLogger } from '../utils/debugLogger.js';
 import type { AgentLoopContext } from '../config/agent-loop-context.js';
+import type { HistoryTurn } from '../core/agentChatHistory.js';
 import {
   SESSION_FILE_PREFIX,
   type TokensSummary,
@@ -497,9 +496,10 @@ export class ChatRecordingService {
     type: ConversationRecordExtra['type'],
     content: PartListUnion,
     displayContent?: PartListUnion,
+    id?: string,
   ): MessageRecord {
     return {
-      id: randomUUID(),
+      id: id || randomUUID(),
       timestamp: new Date().toISOString(),
       type,
       content,
@@ -512,14 +512,17 @@ export class ChatRecordingService {
     type: ConversationRecordExtra['type'];
     content: PartListUnion;
     displayContent?: PartListUnion;
-  }): void {
-    if (!this.conversationFile || !this.cachedConversation) return;
+    id?: string;
+  }): string {
+    if (!this.conversationFile || !this.cachedConversation)
+      return message.id || randomUUID();
 
     try {
       const msg = this.newMessage(
         message.type,
         message.content,
         message.displayContent,
+        message.id,
       );
       if (msg.type === 'gemini') {
         msg.thoughts = this.queuedThoughts;
@@ -530,12 +533,30 @@ export class ChatRecordingService {
       }
       this.pushMessage(msg);
       this.updateMetadata({ lastUpdated: new Date().toISOString() });
+      return msg.id;
     } catch (error) {
       debugLogger.error('Error saving message to chat history.', error);
       throw error;
     }
   }
 
+  /**
+   * Records a synthetic message (e.g. Binary Received, Snapshot/Summary)
+   * and returns its durable ID.
+   */
+  recordSyntheticMessage(
+    type: ConversationRecordExtra['type'],
+    content: PartListUnion,
+    id?: string,
+  ): string {
+    return this.recordMessage({
+      model: undefined,
+      type,
+      content,
+      id,
+    });
+  }
+
   recordThought(thought: ThoughtSummary): void {
     if (!this.conversationFile) return;
     this.queuedThoughts.push({
@@ -869,48 +890,83 @@ export class ChatRecordingService {
     return this.cachedConversation;
   }
 
-  updateMessagesFromHistory(history: readonly Content[]): void {
+  updateMessagesFromHistory(history: readonly HistoryTurn[]): void {
     if (!this.conversationFile || !this.cachedConversation) return;
 
     try {
-      const partsMap = new Map<string, Part[]>();
-      for (const content of history) {
-        if (content.role === 'user' && content.parts) {
-          const callIds = content.parts
-            .map((p) => p.functionResponse?.id)
-            .filter((id): id is string => !!id);
+      let updated = false;
 
-          if (callIds.length === 0) continue;
+      // 1. Sync content and IDs
+      const newMessages: MessageRecord[] = history.map((turn) => {
+        const existing = this.cachedConversation?.messages.find(
+          (m) => m.id === turn.id,
+        );
 
-          let currentCallId = callIds[0];
-          for (const part of content.parts) {
-            if (part.functionResponse?.id) {
-              currentCallId = part.functionResponse.id;
+        if (existing) {
+          // If content parts have changed (e.g. masking), update them
+          if (
+            JSON.stringify(existing.content) !==
+            JSON.stringify(turn.content.parts)
+          ) {
+            updated = true;
+          }
+          return {
+            ...existing,
+            content: turn.content.parts || [],
+          };
+        }
+
+        // It's a new (possibly synthetic) turn like a summary
+        updated = true;
+        return this.newMessage(
+          turn.content.role === 'user' ? 'user' : 'gemini',
+          turn.content.parts || [],
+          undefined,
+          turn.id,
+        );
+      });
+
+      // 2. Specialized 'Masking Sync' for tool call results
+      // If a user turn in history contains a functionResponse, we update the
+      // corresponding ToolCallRecord in the preceding gemini message.
+      for (const turn of history) {
+        if (turn.content.role !== 'user') continue;
+        for (const part of turn.content.parts || []) {
+          if (part.functionResponse) {
+            const callId = part.functionResponse.id;
+            // Find the gemini message that contains this tool call
+            const geminiMsg = newMessages.find(
+              (m) =>
+                m.type === 'gemini' &&
+                m.toolCalls?.some((tc) => tc.id === callId),
+            );
+            if (geminiMsg && geminiMsg.type === 'gemini') {
+              const tc = geminiMsg.toolCalls!.find((tc) => tc.id === callId);
+              if (tc) {
+                // If the history version is different (e.g. masked), sync it into the record
+                // We sync the entire parts array of the user turn to ensure sibling parts are preserved
+                if (
+                  JSON.stringify(tc.result) !==
+                  JSON.stringify(turn.content.parts)
+                ) {
+                  tc.result = turn.content.parts || [];
+                  updated = true;
+                }
+              }
             }
-
-            if (!partsMap.has(currentCallId)) {
-              partsMap.set(currentCallId, []);
-            }
-            partsMap.get(currentCallId)!.push(part);
           }
         }
       }
 
-      for (const message of this.cachedConversation.messages) {
-        let msgChanged = false;
-        if (message.type === 'gemini' && message.toolCalls) {
-          for (const toolCall of message.toolCalls) {
-            const newParts = partsMap.get(toolCall.id);
-            if (newParts !== undefined) {
-              toolCall.result = newParts;
-              msgChanged = true;
-            }
-          }
-        }
-        if (msgChanged) {
-          // Push updated message to log
-          this.pushMessage(message);
-        }
+      if (
+        updated ||
+        newMessages.length !== this.cachedConversation.messages.length
+      ) {
+        this.cachedConversation.messages = newMessages;
+        this.updateMetadata({
+          messages: newMessages,
+          lastUpdated: new Date().toISOString(),
+        });
       }
     } catch (error) {
       debugLogger.error(