From 9455ecd78c3f0990b6ee3dc7955eb7ffea569371 Mon Sep 17 00:00:00 2001 From: Sandy Tao Date: Fri, 6 Mar 2026 19:45:36 -0800 Subject: [PATCH] fix(core): add in-memory cache to ChatRecordingService to prevent OOM (#21502) --- .../src/services/chatRecordingService.test.ts | 124 ++++++++++++++++++ .../core/src/services/chatRecordingService.ts | 76 +++++++---- 2 files changed, 177 insertions(+), 23 deletions(-) diff --git a/packages/core/src/services/chatRecordingService.test.ts b/packages/core/src/services/chatRecordingService.test.ts index 5aaa0a2538..2b8e8f1977 100644 --- a/packages/core/src/services/chatRecordingService.test.ts +++ b/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', () => { diff --git a/packages/core/src/services/chatRecordingService.ts b/packages/core/src/services/chatRecordingService.ts index 1748ccbe20..cd8d1e53c1 100644 --- a/packages/core/src/services/chatRecordingService.ts +++ b/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 }