From 36ea986cfe443d2d363db6e6daa3a0ced7408f3b Mon Sep 17 00:00:00 2001 From: bl-ue <54780737+bl-ue@users.noreply.github.com> Date: Mon, 18 Aug 2025 18:39:57 -0600 Subject: [PATCH] feat(sessions): Introduce core ChatRecordingService for automatic conversation saving (#5221) --- packages/core/src/config/config.ts | 6 +- packages/core/src/index.ts | 1 + .../src/services/chatRecordingService.test.ts | 367 +++++++++++++++ .../core/src/services/chatRecordingService.ts | 433 ++++++++++++++++++ 4 files changed, 806 insertions(+), 1 deletion(-) create mode 100644 packages/core/src/services/chatRecordingService.test.ts create mode 100644 packages/core/src/services/chatRecordingService.ts diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index ad4f8ed9..751012e7 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -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; } diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts index 82ffa1ef..45f7e4ce 100644 --- a/packages/core/src/index.ts +++ b/packages/core/src/index.ts @@ -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 diff --git a/packages/core/src/services/chatRecordingService.test.ts b/packages/core/src/services/chatRecordingService.test.ts new file mode 100644 index 00000000..b78fdde2 --- /dev/null +++ b/packages/core/src/services/chatRecordingService.test.ts @@ -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; + let writeFileSyncSpy: MockInstance; + + 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', + ); + }); + }); +}); diff --git a/packages/core/src/services/chatRecordingService.ts b/packages/core/src/services/chatRecordingService.ts new file mode 100644 index 00000000..9286fcdf --- /dev/null +++ b/packages/core/src/services/chatRecordingService.ts @@ -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; + 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; + 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//chats/ + */ +export class ChatRecordingService { + private conversationFile: string | null = null; + private cachedLastConvData: string | null = null; + private sessionId: string; + private projectHash: string; + private queuedThoughts: Array = []; + 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; + } + } +}