mirror of
https://github.com/QwenLM/qwen-code.git
synced 2025-12-22 01:37:50 +00:00
# 🚀 Sync Gemini CLI v0.2.1 - Major Feature Update (#483)
This commit is contained in:
tanzhenxin
GitHub
433
packages/core/src/services/chatRecordingService.ts
Normal file
433
packages/core/src/services/chatRecordingService.ts
Normal file
@@ -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;
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user