feat(sessions): Introduce core ChatRecordingService for automatic conversation saving (#5221)
This commit is contained in:
bl-ue
GitHub
parent
6fc68ff8d4
commit
36ea986cfe
|
|
@ -207,7 +207,7 @@ export interface ConfigParameters {
|
|||
export class Config {
|
||||
private toolRegistry!: ToolRegistry;
|
||||
private promptRegistry!: PromptRegistry;
|
||||
private readonly sessionId: string;
|
||||
private sessionId: string;
|
||||
private fileSystemService: FileSystemService;
|
||||
private contentGeneratorConfig!: ContentGeneratorConfig;
|
||||
private readonly embeddingModel: string;
|
||||
|
|
@ -409,6 +409,10 @@ export class Config {
|
|||
return this.sessionId;
|
||||
}
|
||||
|
||||
setSessionId(sessionId: string): void {
|
||||
this.sessionId = sessionId;
|
||||
}
|
||||
|
||||
shouldLoadMemoryFromIncludeDirectories(): boolean {
|
||||
return this.loadMemoryFromIncludeDirectories;
|
||||
}
|
||||
|
|
|
|||
|
|
@ -46,6 +46,7 @@ export * from './utils/errorParsing.js';
|
|||
// Export services
|
||||
export * from './services/fileDiscoveryService.js';
|
||||
export * from './services/gitService.js';
|
||||
export * from './services/chatRecordingService.js';
|
||||
export * from './services/fileSystemService.js';
|
||||
|
||||
// Export IDE specific logic
|
||||
|
|
|
|||
|
|
@ -0,0 +1,367 @@
|
|||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import {
|
||||
expect,
|
||||
it,
|
||||
describe,
|
||||
vi,
|
||||
beforeEach,
|
||||
afterEach,
|
||||
MockInstance,
|
||||
} from 'vitest';
|
||||
import fs from 'node:fs';
|
||||
import path from 'node:path';
|
||||
import { randomUUID } from 'node:crypto';
|
||||
import {
|
||||
ChatRecordingService,
|
||||
ConversationRecord,
|
||||
ToolCallRecord,
|
||||
} from './chatRecordingService.js';
|
||||
import { Config } from '../config/config.js';
|
||||
import { getProjectHash } from '../utils/paths.js';
|
||||
|
||||
vi.mock('node:fs');
|
||||
vi.mock('node:path');
|
||||
vi.mock('node:crypto');
|
||||
vi.mock('../utils/paths.js');
|
||||
|
||||
describe('ChatRecordingService', () => {
|
||||
let chatRecordingService: ChatRecordingService;
|
||||
let mockConfig: Config;
|
||||
|
||||
let mkdirSyncSpy: MockInstance<typeof fs.mkdirSync>;
|
||||
let writeFileSyncSpy: MockInstance<typeof fs.writeFileSync>;
|
||||
|
||||
beforeEach(() => {
|
||||
mockConfig = {
|
||||
getSessionId: vi.fn().mockReturnValue('test-session-id'),
|
||||
getProjectRoot: vi.fn().mockReturnValue('/test/project/root'),
|
||||
getProjectTempDir: vi
|
||||
.fn()
|
||||
.mockReturnValue('/test/project/root/.gemini/tmp'),
|
||||
getModel: vi.fn().mockReturnValue('gemini-pro'),
|
||||
getDebugMode: vi.fn().mockReturnValue(false),
|
||||
} as unknown as Config;
|
||||
|
||||
vi.mocked(getProjectHash).mockReturnValue('test-project-hash');
|
||||
vi.mocked(randomUUID).mockReturnValue('this-is-a-test-uuid');
|
||||
vi.mocked(path.join).mockImplementation((...args) => args.join('/'));
|
||||
|
||||
chatRecordingService = new ChatRecordingService(mockConfig);
|
||||
|
||||
mkdirSyncSpy = vi
|
||||
.spyOn(fs, 'mkdirSync')
|
||||
.mockImplementation(() => undefined);
|
||||
|
||||
writeFileSyncSpy = vi
|
||||
.spyOn(fs, 'writeFileSync')
|
||||
.mockImplementation(() => undefined);
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
vi.restoreAllMocks();
|
||||
});
|
||||
|
||||
describe('initialize', () => {
|
||||
it('should create a new session if none is provided', () => {
|
||||
chatRecordingService.initialize();
|
||||
|
||||
expect(mkdirSyncSpy).toHaveBeenCalledWith(
|
||||
'/test/project/root/.gemini/tmp/chats',
|
||||
{ recursive: true },
|
||||
);
|
||||
expect(writeFileSyncSpy).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should resume from an existing session if provided', () => {
|
||||
const readFileSyncSpy = vi.spyOn(fs, 'readFileSync').mockReturnValue(
|
||||
JSON.stringify({
|
||||
sessionId: 'old-session-id',
|
||||
projectHash: 'test-project-hash',
|
||||
messages: [],
|
||||
}),
|
||||
);
|
||||
const writeFileSyncSpy = vi
|
||||
.spyOn(fs, 'writeFileSync')
|
||||
.mockImplementation(() => undefined);
|
||||
|
||||
chatRecordingService.initialize({
|
||||
filePath: '/test/project/root/.gemini/tmp/chats/session.json',
|
||||
conversation: {
|
||||
sessionId: 'old-session-id',
|
||||
} as ConversationRecord,
|
||||
});
|
||||
|
||||
expect(mkdirSyncSpy).not.toHaveBeenCalled();
|
||||
expect(readFileSyncSpy).toHaveBeenCalled();
|
||||
expect(writeFileSyncSpy).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
||||
describe('recordMessage', () => {
|
||||
beforeEach(() => {
|
||||
chatRecordingService.initialize();
|
||||
vi.spyOn(fs, 'readFileSync').mockReturnValue(
|
||||
JSON.stringify({
|
||||
sessionId: 'test-session-id',
|
||||
projectHash: 'test-project-hash',
|
||||
messages: [],
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it('should record a new message', () => {
|
||||
const writeFileSyncSpy = vi
|
||||
.spyOn(fs, 'writeFileSync')
|
||||
.mockImplementation(() => undefined);
|
||||
chatRecordingService.recordMessage({ type: 'user', content: 'Hello' });
|
||||
expect(mkdirSyncSpy).toHaveBeenCalled();
|
||||
expect(writeFileSyncSpy).toHaveBeenCalled();
|
||||
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');
|
||||
expect(conversation.messages[0].type).toBe('user');
|
||||
});
|
||||
|
||||
it('should append to the last message if append is true and types match', () => {
|
||||
const writeFileSyncSpy = vi
|
||||
.spyOn(fs, 'writeFileSync')
|
||||
.mockImplementation(() => undefined);
|
||||
const initialConversation = {
|
||||
sessionId: 'test-session-id',
|
||||
projectHash: 'test-project-hash',
|
||||
messages: [
|
||||
{
|
||||
id: '1',
|
||||
type: 'user',
|
||||
content: 'Hello',
|
||||
timestamp: new Date().toISOString(),
|
||||
},
|
||||
],
|
||||
};
|
||||
vi.spyOn(fs, 'readFileSync').mockReturnValue(
|
||||
JSON.stringify(initialConversation),
|
||||
);
|
||||
|
||||
chatRecordingService.recordMessage({
|
||||
type: 'user',
|
||||
content: ' World',
|
||||
append: true,
|
||||
});
|
||||
|
||||
expect(mkdirSyncSpy).toHaveBeenCalled();
|
||||
expect(writeFileSyncSpy).toHaveBeenCalled();
|
||||
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');
|
||||
});
|
||||
});
|
||||
|
||||
describe('recordThought', () => {
|
||||
it('should queue a thought', () => {
|
||||
chatRecordingService.initialize();
|
||||
chatRecordingService.recordThought({
|
||||
subject: 'Thinking',
|
||||
description: 'Thinking...',
|
||||
});
|
||||
// @ts-expect-error private property
|
||||
expect(chatRecordingService.queuedThoughts).toHaveLength(1);
|
||||
// @ts-expect-error private property
|
||||
expect(chatRecordingService.queuedThoughts[0].subject).toBe('Thinking');
|
||||
// @ts-expect-error private property
|
||||
expect(chatRecordingService.queuedThoughts[0].description).toBe(
|
||||
'Thinking...',
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe('recordMessageTokens', () => {
|
||||
beforeEach(() => {
|
||||
chatRecordingService.initialize();
|
||||
});
|
||||
|
||||
it('should update the last message with token info', () => {
|
||||
const writeFileSyncSpy = vi
|
||||
.spyOn(fs, 'writeFileSync')
|
||||
.mockImplementation(() => undefined);
|
||||
const initialConversation = {
|
||||
sessionId: 'test-session-id',
|
||||
projectHash: 'test-project-hash',
|
||||
messages: [
|
||||
{
|
||||
id: '1',
|
||||
type: 'gemini',
|
||||
content: 'Response',
|
||||
timestamp: new Date().toISOString(),
|
||||
},
|
||||
],
|
||||
};
|
||||
vi.spyOn(fs, 'readFileSync').mockReturnValue(
|
||||
JSON.stringify(initialConversation),
|
||||
);
|
||||
|
||||
chatRecordingService.recordMessageTokens({
|
||||
input: 1,
|
||||
output: 2,
|
||||
total: 3,
|
||||
cached: 0,
|
||||
});
|
||||
|
||||
expect(mkdirSyncSpy).toHaveBeenCalled();
|
||||
expect(writeFileSyncSpy).toHaveBeenCalled();
|
||||
const conversation = JSON.parse(
|
||||
writeFileSyncSpy.mock.calls[0][1] as string,
|
||||
) as ConversationRecord;
|
||||
expect(conversation.messages[0]).toEqual({
|
||||
...initialConversation.messages[0],
|
||||
tokens: { input: 1, output: 2, total: 3, cached: 0 },
|
||||
});
|
||||
});
|
||||
|
||||
it('should queue token info if the last message already has tokens', () => {
|
||||
const initialConversation = {
|
||||
sessionId: 'test-session-id',
|
||||
projectHash: 'test-project-hash',
|
||||
messages: [
|
||||
{
|
||||
id: '1',
|
||||
type: 'gemini',
|
||||
content: 'Response',
|
||||
timestamp: new Date().toISOString(),
|
||||
tokens: { input: 1, output: 1, total: 2, cached: 0 },
|
||||
},
|
||||
],
|
||||
};
|
||||
vi.spyOn(fs, 'readFileSync').mockReturnValue(
|
||||
JSON.stringify(initialConversation),
|
||||
);
|
||||
|
||||
chatRecordingService.recordMessageTokens({
|
||||
input: 2,
|
||||
output: 2,
|
||||
total: 4,
|
||||
cached: 0,
|
||||
});
|
||||
|
||||
// @ts-expect-error private property
|
||||
expect(chatRecordingService.queuedTokens).toEqual({
|
||||
input: 2,
|
||||
output: 2,
|
||||
total: 4,
|
||||
cached: 0,
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe('recordToolCalls', () => {
|
||||
beforeEach(() => {
|
||||
chatRecordingService.initialize();
|
||||
});
|
||||
|
||||
it('should add new tool calls to the last message', () => {
|
||||
const writeFileSyncSpy = vi
|
||||
.spyOn(fs, 'writeFileSync')
|
||||
.mockImplementation(() => undefined);
|
||||
const initialConversation = {
|
||||
sessionId: 'test-session-id',
|
||||
projectHash: 'test-project-hash',
|
||||
messages: [
|
||||
{
|
||||
id: '1',
|
||||
type: 'gemini',
|
||||
content: '',
|
||||
timestamp: new Date().toISOString(),
|
||||
},
|
||||
],
|
||||
};
|
||||
vi.spyOn(fs, 'readFileSync').mockReturnValue(
|
||||
JSON.stringify(initialConversation),
|
||||
);
|
||||
|
||||
const toolCall: ToolCallRecord = {
|
||||
id: 'tool-1',
|
||||
name: 'testTool',
|
||||
args: {},
|
||||
status: 'awaiting_approval',
|
||||
timestamp: new Date().toISOString(),
|
||||
};
|
||||
chatRecordingService.recordToolCalls([toolCall]);
|
||||
|
||||
expect(mkdirSyncSpy).toHaveBeenCalled();
|
||||
expect(writeFileSyncSpy).toHaveBeenCalled();
|
||||
const conversation = JSON.parse(
|
||||
writeFileSyncSpy.mock.calls[0][1] as string,
|
||||
) as ConversationRecord;
|
||||
expect(conversation.messages[0]).toEqual({
|
||||
...initialConversation.messages[0],
|
||||
toolCalls: [toolCall],
|
||||
});
|
||||
});
|
||||
|
||||
it('should create a new message if the last message is not from gemini', () => {
|
||||
const writeFileSyncSpy = vi
|
||||
.spyOn(fs, 'writeFileSync')
|
||||
.mockImplementation(() => undefined);
|
||||
const initialConversation = {
|
||||
sessionId: 'test-session-id',
|
||||
projectHash: 'test-project-hash',
|
||||
messages: [
|
||||
{
|
||||
id: 'a-uuid',
|
||||
type: 'user',
|
||||
content: 'call a tool',
|
||||
timestamp: new Date().toISOString(),
|
||||
},
|
||||
],
|
||||
};
|
||||
vi.spyOn(fs, 'readFileSync').mockReturnValue(
|
||||
JSON.stringify(initialConversation),
|
||||
);
|
||||
|
||||
const toolCall: ToolCallRecord = {
|
||||
id: 'tool-1',
|
||||
name: 'testTool',
|
||||
args: {},
|
||||
status: 'awaiting_approval',
|
||||
timestamp: new Date().toISOString(),
|
||||
};
|
||||
chatRecordingService.recordToolCalls([toolCall]);
|
||||
|
||||
expect(mkdirSyncSpy).toHaveBeenCalled();
|
||||
expect(writeFileSyncSpy).toHaveBeenCalled();
|
||||
const conversation = JSON.parse(
|
||||
writeFileSyncSpy.mock.calls[0][1] as string,
|
||||
) as ConversationRecord;
|
||||
expect(conversation.messages).toHaveLength(2);
|
||||
expect(conversation.messages[1]).toEqual({
|
||||
...conversation.messages[1],
|
||||
id: 'this-is-a-test-uuid',
|
||||
model: 'gemini-pro',
|
||||
type: 'gemini',
|
||||
thoughts: [],
|
||||
content: '',
|
||||
toolCalls: [toolCall],
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
describe('deleteSession', () => {
|
||||
it('should delete the session file', () => {
|
||||
const unlinkSyncSpy = vi
|
||||
.spyOn(fs, 'unlinkSync')
|
||||
.mockImplementation(() => undefined);
|
||||
chatRecordingService.deleteSession('test-session-id');
|
||||
expect(unlinkSyncSpy).toHaveBeenCalledWith(
|
||||
'/test/project/root/.gemini/tmp/chats/test-session-id.json',
|
||||
);
|
||||
});
|
||||
});
|
||||
});
|
||||
|
|
@ -0,0 +1,433 @@
|
|||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { type Config } from '../config/config.js';
|
||||
import { type Status } from '../core/coreToolScheduler.js';
|
||||
import { type ThoughtSummary } from '../core/turn.js';
|
||||
import { getProjectHash } from '../utils/paths.js';
|
||||
import path from 'node:path';
|
||||
import fs from 'node:fs';
|
||||
import { randomUUID } from 'node:crypto';
|
||||
import { PartListUnion } from '@google/genai';
|
||||
|
||||
/**
|
||||
* Token usage summary for a message or conversation.
|
||||
*/
|
||||
export interface TokensSummary {
|
||||
input: number; // promptTokenCount
|
||||
output: number; // candidatesTokenCount
|
||||
cached: number; // cachedContentTokenCount
|
||||
thoughts?: number; // thoughtsTokenCount
|
||||
tool?: number; // toolUsePromptTokenCount
|
||||
total: number; // totalTokenCount
|
||||
}
|
||||
|
||||
/**
|
||||
* Base fields common to all messages.
|
||||
*/
|
||||
export interface BaseMessageRecord {
|
||||
id: string;
|
||||
timestamp: string;
|
||||
content: string;
|
||||
}
|
||||
|
||||
/**
|
||||
* Record of a tool call execution within a conversation.
|
||||
*/
|
||||
export interface ToolCallRecord {
|
||||
id: string;
|
||||
name: string;
|
||||
args: Record<string, unknown>;
|
||||
result?: PartListUnion | null;
|
||||
status: Status;
|
||||
timestamp: string;
|
||||
// UI-specific fields for display purposes
|
||||
displayName?: string;
|
||||
description?: string;
|
||||
resultDisplay?: string;
|
||||
renderOutputAsMarkdown?: boolean;
|
||||
}
|
||||
|
||||
/**
|
||||
* Message type and message type-specific fields.
|
||||
*/
|
||||
export type ConversationRecordExtra =
|
||||
| {
|
||||
type: 'user';
|
||||
}
|
||||
| {
|
||||
type: 'gemini';
|
||||
toolCalls?: ToolCallRecord[];
|
||||
thoughts?: Array<ThoughtSummary & { timestamp: string }>;
|
||||
tokens?: TokensSummary | null;
|
||||
model?: string;
|
||||
};
|
||||
|
||||
/**
|
||||
* A single message record in a conversation.
|
||||
*/
|
||||
export type MessageRecord = BaseMessageRecord & ConversationRecordExtra;
|
||||
|
||||
/**
|
||||
* Complete conversation record stored in session files.
|
||||
*/
|
||||
export interface ConversationRecord {
|
||||
sessionId: string;
|
||||
projectHash: string;
|
||||
startTime: string;
|
||||
lastUpdated: string;
|
||||
messages: MessageRecord[];
|
||||
}
|
||||
|
||||
/**
|
||||
* Data structure for resuming an existing session.
|
||||
*/
|
||||
export interface ResumedSessionData {
|
||||
conversation: ConversationRecord;
|
||||
filePath: string;
|
||||
}
|
||||
|
||||
/**
|
||||
* Service for automatically recording chat conversations to disk.
|
||||
*
|
||||
* This service provides comprehensive conversation recording that captures:
|
||||
* - All user and assistant messages
|
||||
* - Tool calls and their execution results
|
||||
* - Token usage statistics
|
||||
* - Assistant thoughts and reasoning
|
||||
*
|
||||
* Sessions are stored as JSON files in ~/.gemini/tmp/<project_hash>/chats/
|
||||
*/
|
||||
export class ChatRecordingService {
|
||||
private conversationFile: string | null = null;
|
||||
private cachedLastConvData: string | null = null;
|
||||
private sessionId: string;
|
||||
private projectHash: string;
|
||||
private queuedThoughts: Array<ThoughtSummary & { timestamp: string }> = [];
|
||||
private queuedTokens: TokensSummary | null = null;
|
||||
private config: Config;
|
||||
|
||||
constructor(config: Config) {
|
||||
this.config = config;
|
||||
this.sessionId = config.getSessionId();
|
||||
this.projectHash = getProjectHash(config.getProjectRoot());
|
||||
}
|
||||
|
||||
/**
|
||||
* Initializes the chat recording service: creates a new conversation file and associates it with
|
||||
* this service instance, or resumes from an existing session if resumedSessionData is provided.
|
||||
*/
|
||||
initialize(resumedSessionData?: ResumedSessionData): void {
|
||||
try {
|
||||
if (resumedSessionData) {
|
||||
// Resume from existing session
|
||||
this.conversationFile = resumedSessionData.filePath;
|
||||
this.sessionId = resumedSessionData.conversation.sessionId;
|
||||
|
||||
// Update the session ID in the existing file
|
||||
this.updateConversation((conversation) => {
|
||||
conversation.sessionId = this.sessionId;
|
||||
});
|
||||
|
||||
// Clear any cached data to force fresh reads
|
||||
this.cachedLastConvData = null;
|
||||
} else {
|
||||
// Create new session
|
||||
const chatsDir = path.join(this.config.getProjectTempDir(), 'chats');
|
||||
fs.mkdirSync(chatsDir, { recursive: true });
|
||||
|
||||
const timestamp = new Date()
|
||||
.toISOString()
|
||||
.slice(0, 16)
|
||||
.replace(/:/g, '-');
|
||||
const filename = `session-${timestamp}-${this.sessionId.slice(
|
||||
0,
|
||||
8,
|
||||
)}.json`;
|
||||
this.conversationFile = path.join(chatsDir, filename);
|
||||
|
||||
this.writeConversation({
|
||||
sessionId: this.sessionId,
|
||||
projectHash: this.projectHash,
|
||||
startTime: new Date().toISOString(),
|
||||
lastUpdated: new Date().toISOString(),
|
||||
messages: [],
|
||||
});
|
||||
}
|
||||
|
||||
// Clear any queued data since this is a fresh start
|
||||
this.queuedThoughts = [];
|
||||
this.queuedTokens = null;
|
||||
} catch (error) {
|
||||
console.error('Error initializing chat recording service:', error);
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
|
||||
private getLastMessage(
|
||||
conversation: ConversationRecord,
|
||||
): MessageRecord | undefined {
|
||||
return conversation.messages.at(-1);
|
||||
}
|
||||
|
||||
private newMessage(
|
||||
type: ConversationRecordExtra['type'],
|
||||
content: string,
|
||||
): MessageRecord {
|
||||
return {
|
||||
id: randomUUID(),
|
||||
timestamp: new Date().toISOString(),
|
||||
type,
|
||||
content,
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Records a message in the conversation.
|
||||
*/
|
||||
recordMessage(message: {
|
||||
type: ConversationRecordExtra['type'];
|
||||
content: string;
|
||||
append?: boolean;
|
||||
}): 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.
|
||||
conversation.messages.push({
|
||||
...msg,
|
||||
thoughts: this.queuedThoughts,
|
||||
tokens: this.queuedTokens,
|
||||
model: this.config.getModel(),
|
||||
});
|
||||
this.queuedThoughts = [];
|
||||
this.queuedTokens = null;
|
||||
} else {
|
||||
// Or else just add it.
|
||||
conversation.messages.push(msg);
|
||||
}
|
||||
});
|
||||
} catch (error) {
|
||||
console.error('Error saving message:', error);
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Records a thought from the assistant's reasoning process.
|
||||
*/
|
||||
recordThought(thought: ThoughtSummary): void {
|
||||
if (!this.conversationFile) return;
|
||||
|
||||
try {
|
||||
this.queuedThoughts.push({
|
||||
...thought,
|
||||
timestamp: new Date().toISOString(),
|
||||
});
|
||||
} catch (error) {
|
||||
if (this.config.getDebugMode()) {
|
||||
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 {
|
||||
if (!this.conversationFile) return;
|
||||
|
||||
try {
|
||||
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;
|
||||
}
|
||||
});
|
||||
} catch (error) {
|
||||
console.error('Error updating message tokens:', error);
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Adds tool calls to the last message in the conversation (which should be by Gemini).
|
||||
*/
|
||||
recordToolCalls(toolCalls: ToolCallRecord[]): void {
|
||||
if (!this.conversationFile) return;
|
||||
|
||||
try {
|
||||
this.updateConversation((conversation) => {
|
||||
const lastMsg = this.getLastMessage(conversation);
|
||||
// If a tool call was made, but the last message isn't from Gemini, it's because Gemini is
|
||||
// calling tools without starting the message with text. So the user submits a prompt, and
|
||||
// Gemini immediately calls a tool (maybe with some thinking first). In that case, create
|
||||
// a new empty Gemini message.
|
||||
// Also if there are any queued thoughts, it means this tool call(s) is from a new Gemini
|
||||
// message--because it's thought some more since we last, if ever, created a new Gemini
|
||||
// message from tool calls, when we dequeued the thoughts.
|
||||
if (
|
||||
!lastMsg ||
|
||||
lastMsg.type !== 'gemini' ||
|
||||
this.queuedThoughts.length > 0
|
||||
) {
|
||||
const newMsg: MessageRecord = {
|
||||
...this.newMessage('gemini' as const, ''),
|
||||
// This isn't strictly necessary, but TypeScript apparently can't
|
||||
// tell that the first parameter to newMessage() becomes the
|
||||
// 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,
|
||||
thoughts: this.queuedThoughts,
|
||||
model: this.config.getModel(),
|
||||
};
|
||||
// If there are any queued thoughts join them to this message.
|
||||
if (this.queuedThoughts.length > 0) {
|
||||
newMsg.thoughts = this.queuedThoughts;
|
||||
this.queuedThoughts = [];
|
||||
}
|
||||
// If there's any queued tokens info join it to this message.
|
||||
if (this.queuedTokens) {
|
||||
newMsg.tokens = this.queuedTokens;
|
||||
this.queuedTokens = null;
|
||||
}
|
||||
conversation.messages.push(newMsg);
|
||||
} else {
|
||||
// The last message is an existing Gemini message that we need to update.
|
||||
|
||||
// Update any existing tool call entries.
|
||||
if (!lastMsg.toolCalls) {
|
||||
lastMsg.toolCalls = [];
|
||||
}
|
||||
lastMsg.toolCalls = lastMsg.toolCalls.map((toolCall) => {
|
||||
// If there are multiple tool calls with the same ID, this will take the first one.
|
||||
const incomingToolCall = toolCalls.find(
|
||||
(tc) => tc.id === toolCall.id,
|
||||
);
|
||||
if (incomingToolCall) {
|
||||
// Merge in the new data to keep preserve thoughts, etc., that were assigned to older
|
||||
// versions of the tool call.
|
||||
return { ...toolCall, ...incomingToolCall };
|
||||
} else {
|
||||
return toolCall;
|
||||
}
|
||||
});
|
||||
|
||||
// Add any new tools calls that aren't in the message yet.
|
||||
for (const toolCall of toolCalls) {
|
||||
const existingToolCall = lastMsg.toolCalls.find(
|
||||
(tc) => tc.id === toolCall.id,
|
||||
);
|
||||
if (!existingToolCall) {
|
||||
lastMsg.toolCalls.push(toolCall);
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
} catch (error) {
|
||||
console.error('Error adding tool call to message:', error);
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Loads up the conversation record from disk.
|
||||
*/
|
||||
private readConversation(): ConversationRecord {
|
||||
try {
|
||||
this.cachedLastConvData = fs.readFileSync(this.conversationFile!, 'utf8');
|
||||
return JSON.parse(this.cachedLastConvData);
|
||||
} catch (error) {
|
||||
if ((error as NodeJS.ErrnoException).code !== 'ENOENT') {
|
||||
console.error('Error reading conversation file:', error);
|
||||
throw error;
|
||||
}
|
||||
|
||||
// Placeholder empty conversation if file doesn't exist.
|
||||
return {
|
||||
sessionId: this.sessionId,
|
||||
projectHash: this.projectHash,
|
||||
startTime: new Date().toISOString(),
|
||||
lastUpdated: new Date().toISOString(),
|
||||
messages: [],
|
||||
};
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Saves the conversation record; overwrites the file.
|
||||
*/
|
||||
private writeConversation(conversation: ConversationRecord): void {
|
||||
try {
|
||||
if (!this.conversationFile) return;
|
||||
// Don't write the file yet until there's at least one message.
|
||||
if (conversation.messages.length === 0) 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;
|
||||
fs.writeFileSync(this.conversationFile, newContent);
|
||||
}
|
||||
} catch (error) {
|
||||
console.error('Error writing conversation file:', error);
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Convenient helper for updating the conversation without file reading and writing and time
|
||||
* updating boilerplate.
|
||||
*/
|
||||
private updateConversation(
|
||||
updateFn: (conversation: ConversationRecord) => void,
|
||||
) {
|
||||
const conversation = this.readConversation();
|
||||
updateFn(conversation);
|
||||
this.writeConversation(conversation);
|
||||
}
|
||||
|
||||
/**
|
||||
* Deletes a session file by session ID.
|
||||
*/
|
||||
deleteSession(sessionId: string): void {
|
||||
try {
|
||||
const chatsDir = path.join(this.config.getProjectTempDir(), 'chats');
|
||||
const sessionPath = path.join(chatsDir, `${sessionId}.json`);
|
||||
fs.unlinkSync(sessionPath);
|
||||
} catch (error) {
|
||||
console.error('Error deleting session:', error);
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
}
|
||||
Loading…
Reference in New Issue