fix(core): add in-memory cache to ChatRecordingService to prevent OOM (#21502)

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

@@ -245,6 +245,97 @@ describe('ChatRecordingService', () => {
         tool: 0,
       });
     });
+
+    it('should not write to disk when queuing tokens (no last gemini message)', () => {
+      const writeFileSyncSpy = vi.spyOn(fs, 'writeFileSync');
+
+      // Clear spy call count after initialize writes the initial file
+      writeFileSyncSpy.mockClear();
+
+      // No gemini message recorded yet, so tokens should only be queued
+      chatRecordingService.recordMessageTokens({
+        promptTokenCount: 5,
+        candidatesTokenCount: 10,
+        totalTokenCount: 15,
+        cachedContentTokenCount: 0,
+      });
+
+      // writeFileSync should NOT have been called since we only queued
+      expect(writeFileSyncSpy).not.toHaveBeenCalled();
+
+      // @ts-expect-error private property
+      expect(chatRecordingService.queuedTokens).toEqual({
+        input: 5,
+        output: 10,
+        total: 15,
+        cached: 0,
+        thoughts: 0,
+        tool: 0,
+      });
+
+      writeFileSyncSpy.mockRestore();
+    });
+
+    it('should not write to disk when queuing tokens (last message already has tokens)', () => {
+      chatRecordingService.recordMessage({
+        type: 'gemini',
+        content: 'Response',
+        model: 'gemini-pro',
+      });
+
+      // First recordMessageTokens updates the message and writes to disk
+      chatRecordingService.recordMessageTokens({
+        promptTokenCount: 1,
+        candidatesTokenCount: 1,
+        totalTokenCount: 2,
+        cachedContentTokenCount: 0,
+      });
+
+      const writeFileSyncSpy = vi.spyOn(fs, 'writeFileSync');
+      writeFileSyncSpy.mockClear();
+
+      // Second call should only queue, NOT write to disk
+      chatRecordingService.recordMessageTokens({
+        promptTokenCount: 2,
+        candidatesTokenCount: 2,
+        totalTokenCount: 4,
+        cachedContentTokenCount: 0,
+      });
+
+      expect(writeFileSyncSpy).not.toHaveBeenCalled();
+      writeFileSyncSpy.mockRestore();
+    });
+
+    it('should use in-memory cache and not re-read from disk on subsequent operations', () => {
+      chatRecordingService.recordMessage({
+        type: 'gemini',
+        content: 'Response',
+        model: 'gemini-pro',
+      });
+
+      const readFileSyncSpy = vi.spyOn(fs, 'readFileSync');
+      readFileSyncSpy.mockClear();
+
+      // These operations should all use the in-memory cache
+      chatRecordingService.recordMessageTokens({
+        promptTokenCount: 1,
+        candidatesTokenCount: 1,
+        totalTokenCount: 2,
+        cachedContentTokenCount: 0,
+      });
+
+      chatRecordingService.recordMessage({
+        type: 'gemini',
+        content: 'Another response',
+        model: 'gemini-pro',
+      });
+
+      chatRecordingService.saveSummary('Test summary');
+
+      // readFileSync should NOT have been called since we use the in-memory cache
+      expect(readFileSyncSpy).not.toHaveBeenCalled();
+      readFileSyncSpy.mockRestore();
+    });
   });
 
   describe('recordToolCalls', () => {
@@ -769,6 +860,39 @@ describe('ChatRecordingService', () => {
       expect(result[0].text).toBe('Prefix metadata or text');
       expect(result[1].functionResponse!.id).toBe(callId);
     });
+
+    it('should not write to disk when no tool calls match', () => {
+      chatRecordingService.recordMessage({
+        type: 'gemini',
+        content: 'Response with no tool calls',
+        model: 'gemini-pro',
+      });
+
+      const writeFileSyncSpy = vi.spyOn(fs, 'writeFileSync');
+      writeFileSyncSpy.mockClear();
+
+      // History with a tool call ID that doesn't exist in the conversation
+      const history: 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(writeFileSyncSpy).not.toHaveBeenCalled();
+      writeFileSyncSpy.mockRestore();
+    });
   });
 
   describe('ENOENT (missing directory) handling', () => {

packages/core/src/services/chatRecordingService.ts

@@ -128,6 +128,7 @@ export interface ResumedSessionData {
 export class ChatRecordingService {
   private conversationFile: string | null = null;
   private cachedLastConvData: string | null = null;
+  private cachedConversation: ConversationRecord | null = null;
   private sessionId: string;
   private projectHash: string;
   private kind?: 'main' | 'subagent';
@@ -167,6 +168,7 @@ export class ChatRecordingService {
 
         // Clear any cached data to force fresh reads
         this.cachedLastConvData = null;
+        this.cachedConversation = null;
       } else {
         // Create new session
         const chatsDir = path.join(
@@ -308,17 +310,19 @@ export class ChatRecordingService {
         tool: respUsageMetadata.toolUsePromptTokenCount ?? 0,
         total: respUsageMetadata.totalTokenCount ?? 0,
       };
-      this.updateConversation((conversation) => {
-        const lastMsg = this.getLastMessage(conversation);
-        // If the last message already has token info, it's because this new token info is for a
-        // new message that hasn't been recorded yet.
-        if (lastMsg && lastMsg.type === 'gemini' && !lastMsg.tokens) {
-          lastMsg.tokens = tokens;
-          this.queuedTokens = null;
-        } else {
-          this.queuedTokens = tokens;
-        }
-      });
+      const conversation = this.readConversation();
+      const lastMsg = this.getLastMessage(conversation);
+      // If the last message already has token info, it's because this new token info is for a
+      // new message that hasn't been recorded yet.
+      if (lastMsg && lastMsg.type === 'gemini' && !lastMsg.tokens) {
+        lastMsg.tokens = tokens;
+        this.queuedTokens = null;
+        this.writeConversation(conversation);
+      } else {
+        // Only queue tokens in memory; no disk I/O needed since the
+        // conversation record itself hasn't changed.
+        this.queuedTokens = tokens;
+      }
     } catch (error) {
       debugLogger.error(
         'Error updating message tokens in chat history.',
@@ -427,12 +431,32 @@ export class ChatRecordingService {
 
   /**
    * Loads up the conversation record from disk.
+   *
+   * NOTE: The returned object is the live in-memory cache reference.
+   * Any mutations to it will be visible to all subsequent reads.
+   * Callers that mutate the result MUST call writeConversation() to
+   * persist the changes to disk.
    */
   private readConversation(): ConversationRecord {
+    if (this.cachedConversation) {
+      return this.cachedConversation;
+    }
     try {
       this.cachedLastConvData = fs.readFileSync(this.conversationFile!, 'utf8');
-      // eslint-disable-next-line @typescript-eslint/no-unsafe-return
-      return JSON.parse(this.cachedLastConvData);
+      // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
+      this.cachedConversation = JSON.parse(this.cachedLastConvData);
+      if (!this.cachedConversation) {
+        // File is corrupt or contains "null". Fallback to an empty conversation.
+        this.cachedConversation = {
+          sessionId: this.sessionId,
+          projectHash: this.projectHash,
+          startTime: new Date().toISOString(),
+          lastUpdated: new Date().toISOString(),
+          messages: [],
+          kind: this.kind,
+        };
+      }
+      return this.cachedConversation;
     } catch (error) {
       // eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
       if ((error as NodeJS.ErrnoException).code !== 'ENOENT') {
@@ -441,7 +465,7 @@ export class ChatRecordingService {
       }
 
       // Placeholder empty conversation if file doesn't exist.
-      return {
+      this.cachedConversation = {
         sessionId: this.sessionId,
         projectHash: this.projectHash,
         startTime: new Date().toISOString(),
@@ -449,6 +473,7 @@ export class ChatRecordingService {
         messages: [],
         kind: this.kind,
       };
+      return this.cachedConversation;
     }
   }
 
@@ -464,15 +489,19 @@ export class ChatRecordingService {
       // Don't write the file yet until there's at least one message.
       if (conversation.messages.length === 0 && !allowEmpty) return;
 
-      // Only write the file if this change would change the file.
-      if (this.cachedLastConvData !== JSON.stringify(conversation, null, 2)) {
-        conversation.lastUpdated = new Date().toISOString();
-        const newContent = JSON.stringify(conversation, null, 2);
-        this.cachedLastConvData = newContent;
-        // Ensure directory exists before writing (handles cases where temp dir was cleaned)
-        fs.mkdirSync(path.dirname(this.conversationFile), { recursive: true });
-        fs.writeFileSync(this.conversationFile, newContent);
-      }
+      const newContent = JSON.stringify(conversation, null, 2);
+      // Skip the disk write if nothing actually changed (e.g.
+      // updateMessagesFromHistory found no matching tool calls to update).
+      // Compare before updating lastUpdated so the timestamp doesn't
+      // cause a false diff.
+      if (this.cachedLastConvData === newContent) return;
+      this.cachedConversation = conversation;
+      conversation.lastUpdated = new Date().toISOString();
+      const contentToWrite = JSON.stringify(conversation, null, 2);
+      this.cachedLastConvData = contentToWrite;
+      // Ensure directory exists before writing (handles cases where temp dir was cleaned)
+      fs.mkdirSync(path.dirname(this.conversationFile), { recursive: true });
+      fs.writeFileSync(this.conversationFile, contentToWrite);
     } catch (error) {
       // Handle disk full (ENOSPC) gracefully - disable recording but allow conversation to continue
       if (
@@ -482,6 +511,7 @@ export class ChatRecordingService {
         (error as NodeJS.ErrnoException).code === 'ENOSPC'
       ) {
         this.conversationFile = null;
+        this.cachedConversation = null;
         debugLogger.warn(ENOSPC_WARNING_MESSAGE);
         return; // Don't throw - allow the conversation to continue
       }