perf(memory): implement cache eviction in ChatRecordingService and fix leak during resets

This commit is contained in:
Gal Zahavi
2026-04-03 02:36:05 +00:00
parent 21ecd3ca9f
commit 57eea87d41
4 changed files with 115 additions and 7 deletions
+8 -4
View File
@@ -162,10 +162,14 @@ too aggressively.
- [x] Capture baseline memory metrics and initial heap snapshots.
2. **Phase 2: Analysis & Implementation**
- Identify the top 3 memory retainers using Chrome DevTools.
- Implement bounded retention (e.g., capping array sizes in memory,
offloading heavy execution logs to the `.gemini/history` temp files).
- Audit React/Ink components for event listener leaks.
- [x] Identify the top 3 memory retainers using Chrome DevTools (Identified
ChatRecordingService and GeminiChat history).
- [x] Implement bounded retention for ChatRecordingService (implemented
memory-based cache eviction and leak prevention during resets).
- [ ] Implement bounded retention for GeminiChat (improve Tool Output Masking
or add hard history bounds).
- [ ] Audit React/Ink components for event listener leaks.
3. **Phase 3: Validation & CI**
- Run E2E tests to ensure behavioral parity.
- Run `npm run preflight`.
@@ -1117,4 +1117,74 @@ describe('ChatRecordingService', () => {
writeFileSyncSpy.mockRestore();
});
});
describe('Memory management (cache eviction)', () => {
beforeEach(() => {
chatRecordingService.initialize();
});
it('should clear in-memory cache when conversation exceeds 50MB', () => {
// 1. Create a large message (> 50MB)
const largeContent = 'A'.repeat(50 * 1024 * 1024 + 1024);
chatRecordingService.recordMessage({
type: 'user',
content: largeContent,
model: 'gemini-pro',
});
// 2. Check private cache properties
// @ts-expect-error private property
expect(chatRecordingService.cachedConversation).toBeNull();
// @ts-expect-error private property
expect(chatRecordingService.cachedLastConvData).toBeNull();
// 3. Subsequent read should reload from disk
const readFileSyncSpy = vi.spyOn(fs, 'readFileSync');
const conversation = chatRecordingService.getConversation();
expect(conversation).not.toBeNull();
expect(conversation!.messages).toHaveLength(1);
expect(readFileSyncSpy).toHaveBeenCalled();
readFileSyncSpy.mockRestore();
});
it('should keep in-memory cache when conversation is small', () => {
// 1. Create a small message
chatRecordingService.recordMessage({
type: 'user',
content: 'Small message',
model: 'gemini-pro',
});
// 2. Check private cache properties
// @ts-expect-error private property
expect(chatRecordingService.cachedConversation).not.toBeNull();
// @ts-expect-error private property
expect(chatRecordingService.cachedLastConvData).not.toBeNull();
// 3. Subsequent read should NOT reload from disk
const readFileSyncSpy = vi.spyOn(fs, 'readFileSync');
const conversation = chatRecordingService.getConversation();
expect(conversation).not.toBeNull();
expect(readFileSyncSpy).not.toHaveBeenCalled();
readFileSyncSpy.mockRestore();
});
it('should verify writeConversation stringification calls', () => {
const stringifySpy = vi.spyOn(JSON, 'stringify');
// Clear calls from initialize
stringifySpy.mockClear();
chatRecordingService.recordMessage({
type: 'user',
content: 'ping',
model: 'm',
});
// It is called twice: once for comparison with cachedLastConvData,
// and once for writing to disk with the updated lastUpdated timestamp.
expect(stringifySpy).toHaveBeenCalledTimes(2);
stringifySpy.mockRestore();
});
});
});
@@ -27,6 +27,13 @@ import type { AgentLoopContext } from '../config/agent-loop-context.js';
export const SESSION_FILE_PREFIX = 'session-';
/**
* Maximum size of the in-memory chat history cache (50MB).
* When the conversation record exceeds this size, it will be cleared from memory
* after being written to disk to bound the memory footprint.
*/
const MAX_CACHE_SIZE_BYTES = 50 * 1024 * 1024;
/**
* Warning message shown when recording is disabled due to disk full.
*/
@@ -128,6 +135,10 @@ export interface ResumedSessionData {
* - Assistant thoughts and reasoning
*
* Sessions are stored as JSON files in ~/.gemini/tmp/<project_hash>/chats/
*
* Memory Optimization: To prevent unbounded memory growth in long-running
* sessions, this service implements a memory-based eviction policy for its
* in-memory JSON cache.
*/
export class ChatRecordingService {
private conversationFile: string | null = null;
@@ -165,12 +176,15 @@ export class ChatRecordingService {
this.sessionId = resumedSessionData.conversation.sessionId;
this.kind = resumedSessionData.conversation.kind;
// Update the session ID in the existing file
// Use the conversation data for a one-time setup if needed.
// We don't cache it permanently here to save memory; it will be reloaded
// if/when needed by readConversation().
this.updateConversation((conversation) => {
conversation.sessionId = this.sessionId;
});
// Clear any cached data to force fresh reads
// Memory Management: Clear the cache after the initial update
// since it might be huge and we want to allow it to be GC'ed.
this.cachedLastConvData = null;
this.cachedConversation = null;
} else {
@@ -527,13 +541,31 @@ export class ChatRecordingService {
// 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);
// Memory Management: If the conversation is large, clear the in-memory cache
// to bound the heap usage. Subsequent reads will reload from disk.
if (contentToWrite.length > MAX_CACHE_SIZE_BYTES) {
debugLogger.debug(
`[ChatRecordingService] Conversation too large (${(
contentToWrite.length /
1024 /
1024
).toFixed(2)}MB). Evicting from memory cache.`,
);
this.cachedConversation = null;
this.cachedLastConvData = null;
} else {
this.cachedConversation = conversation;
this.cachedLastConvData = contentToWrite;
}
} catch (error) {
// Handle disk full (ENOSPC) gracefully - disable recording but allow conversation to continue
if (
@@ -544,6 +576,7 @@ export class ChatRecordingService {
) {
this.conversationFile = null;
this.cachedConversation = null;
this.cachedLastConvData = null;
debugLogger.warn(ENOSPC_WARNING_MESSAGE);
return; // Don't throw - allow the conversation to continue
}
+1
View File
@@ -1 +1,2 @@
2026-04-03: Completed Phase 1. Implemented scripts/simulate-long-session.ts to reproduce memory growth. Identified ChatRecordingService and GeminiChat history as primary growth sources. Captured baseline metrics showing ~180MB growth per 200MB of tool output data.
2026-04-03: Implemented memory-based cache eviction in ChatRecordingService (50MB threshold). Optimized initialization to prevent carrying over large session records in memory across chat resets. Verified with new unit tests.