diff options
Diffstat (limited to 'packages/core/src/services/chatRecordingService.ts')
| -rw-r--r-- | packages/core/src/services/chatRecordingService.ts | 433 |
1 files changed, 433 insertions, 0 deletions
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<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; + } + } +} |
