feat(core): migrate chat recording to JSONL streaming (#23749)

Author: spencer426

packages/cli/src/ui/hooks/useSessionBrowser.test.ts

@@ -11,14 +11,14 @@ import {
   useSessionBrowser,
   convertSessionToHistoryFormats,
 } from './useSessionBrowser.js';
-import * as fs from 'node:fs/promises';
 import path from 'node:path';
 import { getSessionFiles, type SessionInfo } from '../../utils/sessionUtils.js';
 import {
   type Config,
   type ConversationRecord,
   type MessageRecord,
   CoreToolCallStatus,
+  loadConversationRecord,
 } from '@google/gemini-cli-core';
 import {
   coreEvents,
@@ -46,6 +46,7 @@ vi.mock('@google/gemini-cli-core', async (importOriginal) => {
       clear: vi.fn(),
       hydrate: vi.fn(),
     },
+    loadConversationRecord: vi.fn(),
   };
 });
 
@@ -55,7 +56,6 @@ const MOCKED_SESSION_ID = 'test-session-123';
 const MOCKED_CURRENT_SESSION_ID = 'current-session-id';
 
 describe('useSessionBrowser', () => {
-  const mockedFs = vi.mocked(fs);
   const mockedPath = vi.mocked(path);
   const mockedGetSessionFiles = vi.mocked(getSessionFiles);
 
@@ -98,7 +98,7 @@ describe('useSessionBrowser', () => {
       fileName: MOCKED_FILENAME,
     } as SessionInfo;
     mockedGetSessionFiles.mockResolvedValue([mockSession]);
-    mockedFs.readFile.mockResolvedValue(JSON.stringify(mockConversation));
+    vi.mocked(loadConversationRecord).mockResolvedValue(mockConversation);
 
     const { result } = await renderHook(() =>
       useSessionBrowser(mockConfig, mockOnLoadHistory),
@@ -107,9 +107,8 @@ describe('useSessionBrowser', () => {
     await act(async () => {
       await result.current.handleResumeSession(mockSession);
     });
-    expect(mockedFs.readFile).toHaveBeenCalledWith(
+    expect(loadConversationRecord).toHaveBeenCalledWith(
       `${MOCKED_CHATS_DIR}/${MOCKED_FILENAME}`,
-      'utf8',
     );
     expect(mockConfig.setSessionId).toHaveBeenCalledWith(
       'existing-session-456',
@@ -125,7 +124,9 @@ describe('useSessionBrowser', () => {
       id: MOCKED_SESSION_ID,
       fileName: MOCKED_FILENAME,
     } as SessionInfo;
-    mockedFs.readFile.mockRejectedValue(new Error('File not found'));
+    vi.mocked(loadConversationRecord).mockRejectedValue(
+      new Error('File not found'),
+    );
 
     const { result } = await renderHook(() =>
       useSessionBrowser(mockConfig, mockOnLoadHistory),
@@ -149,7 +150,7 @@ describe('useSessionBrowser', () => {
       id: MOCKED_SESSION_ID,
       fileName: MOCKED_FILENAME,
     } as SessionInfo;
-    mockedFs.readFile.mockResolvedValue('invalid json');
+    vi.mocked(loadConversationRecord).mockResolvedValue(null);
 
     const { result } = await renderHook(() =>
       useSessionBrowser(mockConfig, mockOnLoadHistory),

packages/cli/src/ui/hooks/useSessionBrowser.ts

@@ -6,14 +6,13 @@
 
 import { useState, useCallback } from 'react';
 import type { HistoryItemWithoutId } from '../types.js';
-import * as fs from 'node:fs/promises';
 import path from 'node:path';
 import {
   coreEvents,
   convertSessionToClientHistory,
   uiTelemetryService,
+  loadConversationRecord,
   type Config,
-  type ConversationRecord,
   type ResumedSessionData,
 } from '@google/gemini-cli-core';
 import {
@@ -61,10 +60,12 @@ export const useSessionBrowser = (
           const originalFilePath = path.join(chatsDir, fileName);
 
           // Load up the conversation.
-          // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
-          const conversation: ConversationRecord = JSON.parse(
-            await fs.readFile(originalFilePath, 'utf8'),
-          );
+          const conversation = await loadConversationRecord(originalFilePath);
+          if (!conversation) {
+            throw new Error(
+              `Failed to parse conversation from ${originalFilePath}`,
+            );
+          }
 
           // Use the old session's ID to continue it.
           const existingSessionId = conversation.sessionId;

packages/cli/src/utils/sessionUtils.ts

@@ -12,6 +12,7 @@ import {
   type Storage,
   type ConversationRecord,
   type MessageRecord,
+  loadConversationRecord,
 } from '@google/gemini-cli-core';
 import * as fs from 'node:fs/promises';
 import path from 'node:path';
@@ -250,23 +251,27 @@ export const getAllSessionFiles = async (
   try {
     const files = await fs.readdir(chatsDir);
     const sessionFiles = files
-      .filter((f) => f.startsWith(SESSION_FILE_PREFIX) && f.endsWith('.json'))
+      .filter(
+        (f) =>
+          f.startsWith(SESSION_FILE_PREFIX) &&
+          (f.endsWith('.json') || f.endsWith('.jsonl')),
+      )
       .sort(); // Sort by filename, which includes timestamp
 
     const sessionPromises = sessionFiles.map(
       async (file): Promise<SessionFileEntry> => {
         const filePath = path.join(chatsDir, file);
         try {
-          // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
-          const content: ConversationRecord = JSON.parse(
-            await fs.readFile(filePath, 'utf8'),
-          );
+          const content = await loadConversationRecord(filePath, {
+            metadataOnly: !options.includeFullContent,
+          });
+          if (!content) {
+            return { fileName: file, sessionInfo: null };
+          }
 
           // Validate required fields
           if (
             !content.sessionId ||
-            !content.messages ||
-            !Array.isArray(content.messages) ||
             !content.startTime ||
             !content.lastUpdated
           ) {
@@ -275,7 +280,7 @@ export const getAllSessionFiles = async (
           }
 
           // Skip sessions that only contain system messages (info, error, warning)
-          if (!hasUserOrAssistantMessage(content.messages)) {
+          if (!content.hasUserOrAssistantMessage) {
             return { fileName: file, sessionInfo: null };
           }
 
@@ -285,7 +290,9 @@ export const getAllSessionFiles = async (
             return { fileName: file, sessionInfo: null };
           }
 
-          const firstUserMessage = extractFirstUserMessage(content.messages);
+          const firstUserMessage = content.firstUserMessage
+            ? cleanMessage(content.firstUserMessage)
+            : extractFirstUserMessage(content.messages);
           const isCurrentSession = currentSessionId
             ? file.includes(currentSessionId.slice(0, 8))
             : false;
@@ -310,11 +317,11 @@ export const getAllSessionFiles = async (
 
           const sessionInfo: SessionInfo = {
             id: content.sessionId,
-            file: file.replace('.json', ''),
+            file: file.replace(/\.jsonl?$/, ''),
             fileName: file,
             startTime: content.startTime,
             lastUpdated: content.lastUpdated,
-            messageCount: content.messages.length,
+            messageCount: content.messageCount ?? content.messages.length,
             displayName: content.summary
               ? stripUnsafeCharacters(content.summary)
               : firstUserMessage,
@@ -505,10 +512,10 @@ export class SessionSelector {
     const sessionPath = path.join(chatsDir, sessionInfo.fileName);
 
     try {
-      // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
-      const sessionData: ConversationRecord = JSON.parse(
-        await fs.readFile(sessionPath, 'utf8'),
-      );
+      const sessionData = await loadConversationRecord(sessionPath);
+      if (!sessionData) {
+        throw new Error('Failed to load session data');
+      }
 
       const displayInfo = `Session ${sessionInfo.index}: ${sessionInfo.firstUserMessage} (${sessionInfo.messageCount} messages, ${formatRelativeTime(sessionInfo.lastUpdated)})`;
 

packages/core/src/agents/local-executor.test.ts

@@ -141,6 +141,7 @@ vi.mock('../core/geminiChat.js', () => ({
     CHUNK: 'chunk',
   },
   GeminiChat: vi.fn().mockImplementation(() => ({
+    initialize: vi.fn(),
     sendMessageStream: mockSendMessageStream,
     getHistory: vi.fn((_curated?: boolean) => [...mockChatHistory]),
     setHistory: mockSetHistory,
@@ -434,6 +435,7 @@ describe('LocalAgentExecutor', () => {
     MockedGeminiChat.mockImplementation(
       () =>
         ({
+          initialize: vi.fn(),
           sendMessageStream: mockSendMessageStream,
           setSystemInstruction: mockSetSystemInstruction,
           getHistory: vi.fn((_curated?: boolean) => [...mockChatHistory]),

packages/core/src/agents/local-executor.ts

@@ -1021,15 +1021,16 @@ export class LocalAgentExecutor<TOutput extends z.ZodTypeAny> {
       : undefined;
 
     try {
-      return new GeminiChat(
+      const chat = new GeminiChat(
         this.executionContext,
         systemInstruction,
         [{ functionDeclarations: tools }],
         startHistory,
         undefined,
         undefined,
-        'subagent',
       );
+      await chat.initialize(undefined, 'subagent');
+      return chat;
     } catch (e: unknown) {
       await reportError(
         e,

packages/core/src/config/storage.ts

@@ -353,7 +353,9 @@ export class Storage {
     const chatsDir = path.join(this.getProjectTempDir(), 'chats');
     try {
       const files = await fs.promises.readdir(chatsDir);
-      const jsonFiles = files.filter((f) => f.endsWith('.json'));
+      const jsonFiles = files.filter(
+        (f) => f.endsWith('.json') || f.endsWith('.jsonl'),
+      );
 
       const sessions = await Promise.all(
         jsonFiles.map(async (file) => {

packages/core/src/core/client.test.ts

@@ -63,6 +63,10 @@ vi.mock('node:fs', () => {
     writeFileSync: vi.fn((path: string, data: string) => {
       mockFileSystem.set(path, data);
     }),
+    appendFileSync: vi.fn((path: string, data: string) => {
+      const current = mockFileSystem.get(path) || '';
+      mockFileSystem.set(path, current + data);
+    }),
     readFileSync: vi.fn((path: string) => {
       if (mockFileSystem.has(path)) {
         return mockFileSystem.get(path);

packages/core/src/core/client.ts

@@ -378,7 +378,7 @@ export class GeminiClient {
     try {
       const systemMemory = this.config.getSystemInstructionMemory();
       const systemInstruction = getCoreSystemPrompt(this.config, systemMemory);
-      return new GeminiChat(
+      const chat = new GeminiChat(
         this.config,
         systemInstruction,
         tools,
@@ -392,6 +392,8 @@ export class GeminiClient {
           return [{ functionDeclarations: toolDeclarations }];
         },
       );
+      await chat.initialize(resumedSessionData, 'main');
+      return chat;
     } catch (error) {
       await reportError(
         error,

packages/core/src/core/geminiChat.test.ts

@@ -48,6 +48,10 @@ vi.mock('node:fs', () => {
     writeFileSync: vi.fn((path: string, data: string) => {
       mockFileSystem.set(path, data);
     }),
+    appendFileSync: vi.fn((path: string, data: string) => {
+      const current = mockFileSystem.get(path) || '';
+      mockFileSystem.set(path, current + data);
+    }),
     readFileSync: vi.fn((path: string) => {
       if (mockFileSystem.has(path)) {
         return mockFileSystem.get(path);
@@ -1082,8 +1086,10 @@ describe('GeminiChat', () => {
       );
 
       const { default: fs } = await import('node:fs');
-      const writeFileSync = vi.mocked(fs.writeFileSync);
-      const writeCountBefore = writeFileSync.mock.calls.length;
+      const appendFileSync = vi.mocked(fs.appendFileSync);
+      const writeCountBefore = appendFileSync.mock.calls.length;
+
+      await chat.initialize();
 
       const stream = await chat.sendMessageStream(
         { model: 'test-model' },
@@ -1096,17 +1102,19 @@ describe('GeminiChat', () => {
         // consume
       }
 
-      const newWrites = writeFileSync.mock.calls.slice(writeCountBefore);
+      const newWrites = appendFileSync.mock.calls.slice(writeCountBefore);
       expect(newWrites.length).toBeGreaterThan(0);
 
-      const lastWriteData = JSON.parse(
-        newWrites[newWrites.length - 1][1] as string,
-      ) as { messages: Array<{ type: string }> };
+      const geminiWrite = newWrites.find((w) => {
+        try {
+          const data = JSON.parse(w[1] as string);
+          return data.type === 'gemini';
+        } catch {
+          return false;
+        }
+      });
 
-      const geminiMessages = lastWriteData.messages.filter(
-        (m) => m.type === 'gemini',
-      );
-      expect(geminiMessages.length).toBeGreaterThan(0);
+      expect(geminiWrite).toBeDefined();
     });
   });
 

packages/core/src/core/geminiChat.ts

@@ -256,16 +256,21 @@ export class GeminiChat {
     private history: Content[] = [],
     resumedSessionData?: ResumedSessionData,
     private readonly onModelChanged?: (modelId: string) => Promise<Tool[]>,
-    kind: 'main' | 'subagent' = 'main',
   ) {
     validateHistory(history);
     this.chatRecordingService = new ChatRecordingService(context);
-    this.chatRecordingService.initialize(resumedSessionData, kind);
     this.lastPromptTokenCount = estimateTokenCountSync(
       this.history.flatMap((c) => c.parts || []),
     );
   }
 
+  async initialize(
+    resumedSessionData?: ResumedSessionData,
+    kind: 'main' | 'subagent' = 'main',
+  ) {
+    await this.chatRecordingService.initialize(resumedSessionData, kind);
+  }
+
   setSystemInstruction(sysInstr: string) {
     this.systemInstruction = sysInstr;
   }