feat(sessions): Integrate chat recording into GeminiChat (#6721)

packages/cli/src/nonInteractiveCli.test.ts

@@ -24,11 +24,20 @@ vi.mock('./ui/hooks/atCommandProcessor.js');
 vi.mock('@google/gemini-cli-core', async (importOriginal) => {
   const original =
     await importOriginal<typeof import('@google/gemini-cli-core')>();
+
+  class MockChatRecordingService {
+    initialize = vi.fn();
+    recordMessage = vi.fn();
+    recordMessageTokens = vi.fn();
+    recordToolCalls = vi.fn();
+  }
+
   return {
     ...original,
     executeToolCall: vi.fn(),
     shutdownTelemetry: vi.fn(),
     isTelemetrySdkInitialized: vi.fn().mockReturnValue(true),
+    ChatRecordingService: MockChatRecordingService,
   };
 });
 
@@ -41,6 +50,7 @@ describe('runNonInteractive', () => {
   let processStdoutSpy: vi.SpyInstance;
   let mockGeminiClient: {
     sendMessageStream: vi.Mock;
+    getChatRecordingService: vi.Mock;
   };
 
   beforeEach(async () => {
@@ -59,6 +69,12 @@ describe('runNonInteractive', () => {
 
     mockGeminiClient = {
       sendMessageStream: vi.fn(),
+      getChatRecordingService: vi.fn(() => ({
+        initialize: vi.fn(),
+        recordMessage: vi.fn(),
+        recordMessageTokens: vi.fn(),
+        recordToolCalls: vi.fn(),
+      })),
     };
 
     mockConfig = {
@@ -66,6 +82,11 @@ describe('runNonInteractive', () => {
       getGeminiClient: vi.fn().mockReturnValue(mockGeminiClient),
       getToolRegistry: vi.fn().mockReturnValue(mockToolRegistry),
       getMaxSessionTurns: vi.fn().mockReturnValue(10),
+      getSessionId: vi.fn().mockReturnValue('test-session-id'),
+      getProjectRoot: vi.fn().mockReturnValue('/test/project'),
+      storage: {
+        getProjectTempDir: vi.fn().mockReturnValue('/test/project/.gemini/tmp'),
+      },
       getIdeMode: vi.fn().mockReturnValue(false),
       getFullContext: vi.fn().mockReturnValue(false),
       getContentGeneratorConfig: vi.fn().mockReturnValue({}),
@@ -97,6 +118,10 @@ describe('runNonInteractive', () => {
     const events: ServerGeminiStreamEvent[] = [
       { type: GeminiEventType.Content, value: 'Hello' },
       { type: GeminiEventType.Content, value: ' World' },
+      {
+        type: GeminiEventType.Finished,
+        value: { reason: undefined, usageMetadata: { totalTokenCount: 10 } },
+      },
     ];
     mockGeminiClient.sendMessageStream.mockReturnValue(
       createStreamFromEvents(events),
@@ -132,6 +157,10 @@ describe('runNonInteractive', () => {
     const firstCallEvents: ServerGeminiStreamEvent[] = [toolCallEvent];
     const secondCallEvents: ServerGeminiStreamEvent[] = [
       { type: GeminiEventType.Content, value: 'Final answer' },
+      {
+        type: GeminiEventType.Finished,
+        value: { reason: undefined, usageMetadata: { totalTokenCount: 10 } },
+      },
     ];
 
     mockGeminiClient.sendMessageStream
@@ -187,6 +216,10 @@ describe('runNonInteractive', () => {
         type: GeminiEventType.Content,
         value: 'Sorry, let me try again.',
       },
+      {
+        type: GeminiEventType.Finished,
+        value: { reason: undefined, usageMetadata: { totalTokenCount: 10 } },
+      },
     ];
     mockGeminiClient.sendMessageStream
       .mockReturnValueOnce(createStreamFromEvents([toolCallEvent]))
@@ -242,12 +275,17 @@ describe('runNonInteractive', () => {
     mockCoreExecuteToolCall.mockResolvedValue({
       error: new Error('Tool "nonexistentTool" not found in registry.'),
       resultDisplay: 'Tool "nonexistentTool" not found in registry.',
+      responseParts: [],
     });
     const finalResponse: ServerGeminiStreamEvent[] = [
       {
         type: GeminiEventType.Content,
         value: "Sorry, I can't find that tool.",
       },
+      {
+        type: GeminiEventType.Finished,
+        value: { reason: undefined, usageMetadata: { totalTokenCount: 10 } },
+      },
     ];
 
     mockGeminiClient.sendMessageStream
@@ -304,6 +342,10 @@ describe('runNonInteractive', () => {
     // Mock a simple stream response from the Gemini client
     const events: ServerGeminiStreamEvent[] = [
       { type: GeminiEventType.Content, value: 'Summary complete.' },
+      {
+        type: GeminiEventType.Finished,
+        value: { reason: undefined, usageMetadata: { totalTokenCount: 10 } },
+      },
     ];
     mockGeminiClient.sendMessageStream.mockReturnValue(
       createStreamFromEvents(events),

packages/cli/src/ui/App.test.tsx

@@ -157,7 +157,22 @@ vi.mock('@google/gemini-cli-core', async (importOriginal) => {
         getProjectRoot: vi.fn(() => opts.targetDir),
         getEnablePromptCompletion: vi.fn(() => false),
         getGeminiClient: vi.fn(() => ({
+          isInitialized: vi.fn(() => true),
           getUserTier: vi.fn(),
+          getChatRecordingService: vi.fn(() => ({
+            initialize: vi.fn(),
+            recordMessage: vi.fn(),
+            recordMessageTokens: vi.fn(),
+            recordToolCalls: vi.fn(),
+          })),
+          getChat: vi.fn(() => ({
+            getChatRecordingService: vi.fn(() => ({
+              initialize: vi.fn(),
+              recordMessage: vi.fn(),
+              recordMessageTokens: vi.fn(),
+              recordToolCalls: vi.fn(),
+            })),
+          })),
         })),
         getCheckpointingEnabled: vi.fn(() => opts.checkpointing ?? true),
         getAllGeminiMdFilenames: vi.fn(() => ['GEMINI.md']),

packages/cli/src/ui/hooks/useGeminiStream.test.tsx

@@ -48,6 +48,14 @@ const MockedGeminiClientClass = vi.hoisted(() =>
     this.startChat = mockStartChat;
     this.sendMessageStream = mockSendMessageStream;
     this.addHistory = vi.fn();
+    this.getChatRecordingService = vi.fn().mockReturnValue({
+      recordThought: vi.fn(),
+      initialize: vi.fn(),
+      recordMessage: vi.fn(),
+      recordMessageTokens: vi.fn(),
+      recordToolCalls: vi.fn(),
+      getConversationFile: vi.fn(),
+    });
   }),
 );
 
@@ -1275,7 +1283,10 @@ describe('useGeminiStream', () => {
             type: ServerGeminiEventType.Content,
             value: 'This is a truncated response...',
           };
-          yield { type: ServerGeminiEventType.Finished, value: 'MAX_TOKENS' };
+          yield {
+            type: ServerGeminiEventType.Finished,
+            value: { reason: 'MAX_TOKENS', usageMetadata: undefined },
+          };
         })(),
       );
 
@@ -1324,7 +1335,10 @@ describe('useGeminiStream', () => {
             type: ServerGeminiEventType.Content,
             value: 'Complete response',
           };
-          yield { type: ServerGeminiEventType.Finished, value: 'STOP' };
+          yield {
+            type: ServerGeminiEventType.Finished,
+            value: { reason: 'STOP', usageMetadata: undefined },
+          };
         })(),
       );
 
@@ -1373,7 +1387,10 @@ describe('useGeminiStream', () => {
           };
           yield {
             type: ServerGeminiEventType.Finished,
-            value: 'FINISH_REASON_UNSPECIFIED',
+            value: {
+              reason: 'FINISH_REASON_UNSPECIFIED',
+              usageMetadata: undefined,
+            },
           };
         })(),
       );
@@ -1464,7 +1481,10 @@ describe('useGeminiStream', () => {
               type: ServerGeminiEventType.Content,
               value: `Response for ${reason}`,
             };
-            yield { type: ServerGeminiEventType.Finished, value: reason };
+            yield {
+              type: ServerGeminiEventType.Finished,
+              value: { reason, usageMetadata: undefined },
+            };
           })(),
         );
 
@@ -1579,7 +1599,10 @@ describe('useGeminiStream', () => {
             type: ServerGeminiEventType.Content,
             value: 'Some response content',
           };
-          yield { type: ServerGeminiEventType.Finished, value: 'STOP' };
+          yield {
+            type: ServerGeminiEventType.Finished,
+            value: { reason: 'STOP', usageMetadata: undefined },
+          };
         })(),
       );
 
@@ -1626,7 +1649,10 @@ describe('useGeminiStream', () => {
             type: ServerGeminiEventType.Content,
             value: 'New response content',
           };
-          yield { type: ServerGeminiEventType.Finished, value: 'STOP' };
+          yield {
+            type: ServerGeminiEventType.Finished,
+            value: { reason: 'STOP', usageMetadata: undefined },
+          };
         })(),
       );
 

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

@@ -516,7 +516,10 @@ export const useGeminiStream = (
 
   const handleFinishedEvent = useCallback(
     (event: ServerGeminiFinishedEvent, userMessageTimestamp: number) => {
-      const finishReason = event.value;
+      const finishReason = event.value.reason;
+      if (!finishReason) {
+        return;
+      }
 
       const finishReasonMessages: Record<FinishReason, string | undefined> = {
         [FinishReason.FINISH_REASON_UNSPECIFIED]: undefined,

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

@@ -59,6 +59,7 @@ const mockConfig = {
     model: 'test-model',
     authType: 'oauth-personal',
   }),
+  getGeminiClient: () => null, // No client needed for these tests
 } as unknown as Config;
 
 const mockTool = new MockTool({