MediaMetz/gemini-cli
1
0
Fork 0
mirror of https://github.com/google-gemini/gemini-cli.git synced 2026-04-21 18:44:30 -07:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
bl-ue
2025-09-02 23:29:07 -06:00
committed by GitHub
parent c9bd3ecf6a
commit b5dd6f9ea6
17 changed files with 515 additions and 72 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/core/src/services/chatRecordingService.test.ts
+54 -17
View File
@@ -19,7 +19,14 @@ import { getProjectHash } from '../utils/paths.js';
vi.mock('node:fs');
vi.mock('node:path');
vi.mock('node:crypto');
vi.mock('node:crypto', () => ({
randomUUID: vi.fn(),
createHash: vi.fn(() => ({
update: vi.fn(() => ({
digest: vi.fn(() => 'mocked-hash'),
})),
})),
}));
vi.mock('../utils/paths.js');
describe('ChatRecordingService', () => {
@@ -40,6 +47,13 @@ describe('ChatRecordingService', () => {
},
getModel: vi.fn().mockReturnValue('gemini-pro'),
getDebugMode: vi.fn().mockReturnValue(false),
getToolRegistry: vi.fn().mockReturnValue({
getTool: vi.fn().mockReturnValue({
displayName: 'Test Tool',
description: 'A test tool',
isOutputMarkdown: false,
}),
}),
} as unknown as Config;
vi.mocked(getProjectHash).mockReturnValue('test-project-hash');
@@ -124,7 +138,7 @@ describe('ChatRecordingService', () => {
expect(conversation.messages[0].type).toBe('user');
});
it('should append to the last message if append is true and types match', () => {
it('should create separate messages when recording multiple messages', () => {
const writeFileSyncSpy = vi
.spyOn(fs, 'writeFileSync')
.mockImplementation(() => undefined);
@@ -146,8 +160,7 @@ describe('ChatRecordingService', () => {
chatRecordingService.recordMessage({
type: 'user',
content: ' World',
append: true,
content: 'World',
});
expect(mkdirSyncSpy).toHaveBeenCalled();
@@ -155,8 +168,9 @@ describe('ChatRecordingService', () => {
const conversation = JSON.parse(
writeFileSyncSpy.mock.calls[0][1] as string,
) as ConversationRecord;
expect(conversation.messages).toHaveLength(1);
expect(conversation.messages[0].content).toBe('Hello World');
expect(conversation.messages).toHaveLength(2);
expect(conversation.messages[0].content).toBe('Hello');
expect(conversation.messages[1].content).toBe('World');
});
});
@@ -204,10 +218,10 @@ describe('ChatRecordingService', () => {
);
chatRecordingService.recordMessageTokens({
input: 1,
output: 2,
total: 3,
cached: 0,
promptTokenCount: 1,
candidatesTokenCount: 2,
totalTokenCount: 3,
cachedContentTokenCount: 0,
});
expect(mkdirSyncSpy).toHaveBeenCalled();
@@ -217,7 +231,14 @@ describe('ChatRecordingService', () => {
) as ConversationRecord;
expect(conversation.messages[0]).toEqual({
...initialConversation.messages[0],
tokens: { input: 1, output: 2, total: 3, cached: 0 },
tokens: {
input: 1,
output: 2,
total: 3,
cached: 0,
thoughts: 0,
tool: 0,
},
});
});
@@ -240,10 +261,10 @@ describe('ChatRecordingService', () => {
);
chatRecordingService.recordMessageTokens({
input: 2,
output: 2,
total: 4,
cached: 0,
promptTokenCount: 2,
candidatesTokenCount: 2,
totalTokenCount: 4,
cachedContentTokenCount: 0,
});
// @ts-expect-error private property
@@ -252,6 +273,8 @@ describe('ChatRecordingService', () => {
output: 2,
total: 4,
cached: 0,
thoughts: 0,
tool: 0,
});
});
});
@@ -297,7 +320,14 @@ describe('ChatRecordingService', () => {
) as ConversationRecord;
expect(conversation.messages[0]).toEqual({
...initialConversation.messages[0],
toolCalls: [toolCall],
toolCalls: [
{
...toolCall,
displayName: 'Test Tool',
description: 'A test tool',
renderOutputAsMarkdown: false,
},
],
});
});
@@ -343,7 +373,14 @@ describe('ChatRecordingService', () => {
type: 'gemini',
thoughts: [],
content: '',
toolCalls: [toolCall],
toolCalls: [
{
...toolCall,
displayName: 'Test Tool',
description: 'A test tool',
renderOutputAsMarkdown: false,
},
],
});
});
});
packages/core/src/services/chatRecordingService.ts
+35 -28
View File
@@ -11,7 +11,10 @@ import { getProjectHash } from '../utils/paths.js';
import path from 'node:path';
import fs from 'node:fs';
import { randomUUID } from 'node:crypto';
import type { PartListUnion } from '@google/genai';
import type {
PartListUnion,
GenerateContentResponseUsageMetadata,
} from '@google/genai';
/**
* Token usage summary for a message or conversation.
@@ -31,7 +34,7 @@ export interface TokensSummary {
export interface BaseMessageRecord {
id: string;
timestamp: string;
content: string;
content: PartListUnion;
}
/**
@@ -178,7 +181,7 @@ export class ChatRecordingService {
private newMessage(
type: ConversationRecordExtra['type'],
content: string,
content: PartListUnion,
): MessageRecord {
return {
id: randomUUID(),
@@ -193,22 +196,12 @@ export class ChatRecordingService {
*/
recordMessage(message: {
type: ConversationRecordExtra['type'];
content: string;
append?: boolean;
content: PartListUnion;
}): void {
if (!this.conversationFile) return;
try {
this.updateConversation((conversation) => {
if (message.append) {
const lastMsg = this.getLastMessage(conversation);
if (lastMsg && lastMsg.type === message.type) {
lastMsg.content += message.content;
return;
}
}
// We're not appending, or we are appending but the last message's type is not the same as
// the specified type, so just create a new message.
const msg = this.newMessage(message.type, message.content);
if (msg.type === 'gemini') {
// If it's a new Gemini message then incorporate any queued thoughts.
@@ -243,27 +236,28 @@ export class ChatRecordingService {
timestamp: new Date().toISOString(),
});
} catch (error) {
if (this.config.getDebugMode()) {
console.error('Error saving thought:', error);
throw error;
}
console.error('Error saving thought:', error);
throw error;
}
}
/**
* Updates the tokens for the last message in the conversation (which should be by Gemini).
*/
recordMessageTokens(tokens: {
input: number;
output: number;
cached: number;
thoughts?: number;
tool?: number;
total: number;
}): void {
recordMessageTokens(
respUsageMetadata: GenerateContentResponseUsageMetadata,
): void {
if (!this.conversationFile) return;
try {
const tokens = {
input: respUsageMetadata.promptTokenCount ?? 0,
output: respUsageMetadata.candidatesTokenCount ?? 0,
cached: respUsageMetadata.cachedContentTokenCount ?? 0,
thoughts: respUsageMetadata.thoughtsTokenCount ?? 0,
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
@@ -283,10 +277,23 @@ export class ChatRecordingService {
/**
* Adds tool calls to the last message in the conversation (which should be by Gemini).
* This method enriches tool calls with metadata from the ToolRegistry.
*/
recordToolCalls(toolCalls: ToolCallRecord[]): void {
if (!this.conversationFile) return;
// Enrich tool calls with metadata from the ToolRegistry
const toolRegistry = this.config.getToolRegistry();
const enrichedToolCalls = toolCalls.map((toolCall) => {
const toolInstance = toolRegistry.getTool(toolCall.name);
return {
...toolCall,
displayName: toolInstance?.displayName || toolCall.name,
description: toolInstance?.description || '',
renderOutputAsMarkdown: toolInstance?.isOutputMarkdown || false,
};
});
try {
this.updateConversation((conversation) => {
const lastMsg = this.getLastMessage(conversation);
@@ -309,7 +316,7 @@ export class ChatRecordingService {
// resulting message's type, and so it thinks that toolCalls may
// not be present. Confirming the type here satisfies it.
type: 'gemini' as const,
toolCalls,
toolCalls: enrichedToolCalls,
thoughts: this.queuedThoughts,
model: this.config.getModel(),
};
@@ -346,7 +353,7 @@ export class ChatRecordingService {
});
// Add any new tools calls that aren't in the message yet.
for (const toolCall of toolCalls) {
for (const toolCall of enrichedToolCalls) {
const existingToolCall = lastMsg.toolCalls.find(
(tc) => tc.id === toolCall.id,
);