refactor: Centralize session ID generation and propagation

This commit is contained in:
jerop 2025-06-11 04:46:39 +00:00 committed by Jerop Kipruto
parent 95fdc66e7d
commit d1e23b7c71
20 changed files with 96 additions and 71 deletions

View File

@ -120,6 +120,7 @@ export async function loadCliConfig(
settings: Settings,
extensions: ExtensionConfig[],
geminiIgnorePatterns: string[],
sessionId: string,
): Promise<Config> {
loadEnvironment();
@ -148,6 +149,7 @@ export async function loadCliConfig(
const mcpServers = mergeMcpServers(settings, extensions);
return new Config({
sessionId,
contentGeneratorConfig,
embeddingModel: DEFAULT_GEMINI_EMBEDDING_MODEL,
sandbox: argv.sandbox ?? settings.sandbox,

View File

@ -31,6 +31,7 @@ import {
WebFetchTool,
WebSearchTool,
WriteFileTool,
sessionId,
} from '@gemini-cli/core';
export async function main() {
@ -57,6 +58,7 @@ export async function main() {
settings.merged,
extensions,
geminiIgnorePatterns,
sessionId,
);
// Initialize centralized FileDiscoveryService
@ -180,5 +182,6 @@ async function loadNonInteractiveConfig(
nonInteractiveSettings,
extensions,
config.getGeminiIgnorePatterns(),
config.getSessionId(),
);
}

View File

@ -190,6 +190,7 @@ describe('App UI', () => {
userMemory: '',
geminiMdFileCount: 0,
showMemoryUsage: false,
sessionId: 'test-session-id',
// Provide other required fields for ConfigParameters if necessary
}) as unknown as MockServerConfig;

View File

@ -493,7 +493,7 @@ Add any other context about the problem here.
description: 'save conversation checkpoint. Usage: /save [tag]',
action: async (_mainCommand, subCommand, _args) => {
const tag = (subCommand || '').trim();
const logger = new Logger();
const logger = new Logger(config?.getSessionId() || '');
await logger.initialize();
const chat = await config?.getGeminiClient()?.getChat();
const history = chat?.getHistory() || [];
@ -519,7 +519,7 @@ Add any other context about the problem here.
'resume from conversation checkpoint. Usage: /resume [tag]',
action: async (_mainCommand, subCommand, _args) => {
const tag = (subCommand || '').trim();
const logger = new Logger();
const logger = new Logger(config?.getSessionId() || '');
await logger.initialize();
const conversation = await logger.loadCheckpoint(tag);
if (conversation.length === 0) {

View File

@ -5,6 +5,7 @@
*/
import { useState, useEffect } from 'react';
import { sessionId } from '@gemini-cli/core';
import { Logger } from '@gemini-cli/core';
/**
@ -14,7 +15,7 @@ export const useLogger = () => {
const [logger, setLogger] = useState<Logger | null>(null);
useEffect(() => {
const newLogger = new Logger();
const newLogger = new Logger(sessionId);
/**
* Start async initialization, no need to await. Using await slows down the
* time from launch to see the gemini-cli prompt and it's better to not save

View File

@ -49,6 +49,7 @@ describe('Server Config (config.ts)', () => {
const USER_MEMORY = 'Test User Memory';
const TELEMETRY = false;
const EMBEDDING_MODEL = 'gemini-embedding';
const SESSION_ID = 'test-session-id';
const baseParams: ConfigParameters = {
contentGeneratorConfig: {
apiKey: API_KEY,
@ -62,6 +63,7 @@ describe('Server Config (config.ts)', () => {
fullContext: FULL_CONTEXT,
userMemory: USER_MEMORY,
telemetry: TELEMETRY,
sessionId: SESSION_ID,
};
beforeEach(() => {

View File

@ -55,6 +55,7 @@ export class MCPServerConfig {
}
export interface ConfigParameters {
sessionId: string;
contentGeneratorConfig: ContentGeneratorConfig;
embeddingModel: string;
sandbox?: boolean | string;
@ -83,6 +84,7 @@ export interface ConfigParameters {
export class Config {
private toolRegistry: Promise<ToolRegistry>;
private readonly sessionId: string;
private readonly contentGeneratorConfig: ContentGeneratorConfig;
private readonly embeddingModel: string;
private readonly sandbox: boolean | string | undefined;
@ -111,6 +113,7 @@ export class Config {
private fileDiscoveryService: FileDiscoveryService | null = null;
constructor(params: ConfigParameters) {
this.sessionId = params.sessionId;
this.contentGeneratorConfig = params.contentGeneratorConfig;
this.embeddingModel = params.embeddingModel;
this.sandbox = params.sandbox;
@ -155,6 +158,10 @@ export class Config {
}
}
getSessionId(): string {
return this.sessionId;
}
getContentGeneratorConfig(): ContentGeneratorConfig {
return this.contentGeneratorConfig;
}

View File

@ -108,6 +108,7 @@ describe('Gemini Client (client.ts)', () => {
getUserAgent: vi.fn().mockReturnValue('test-agent'),
getUserMemory: vi.fn().mockReturnValue(''),
getFullContext: vi.fn().mockReturnValue(false),
getSessionId: vi.fn().mockReturnValue('test-session-id'),
};
// eslint-disable-next-line @typescript-eslint/no-explicit-any
return mock as any;

View File

@ -162,6 +162,7 @@ export class GeminiClient {
return new GeminiChat(
await this.contentGenerator,
this.model,
this.config.getSessionId(),
{
systemInstruction,
...this.generateContentConfig,

View File

@ -21,11 +21,12 @@ describe('GeminiChat', () => {
let chat: GeminiChat;
const model = 'gemini-pro';
const config: GenerateContentConfig = {};
const sessionId = 'test-session-id';
beforeEach(() => {
vi.clearAllMocks();
// Reset history for each test by creating a new instance
chat = new GeminiChat(mockModelsModule, model, config, []);
chat = new GeminiChat(mockModelsModule, model, sessionId, config, []);
});
afterEach(() => {
@ -120,7 +121,7 @@ describe('GeminiChat', () => {
chat.recordHistory(userInput, newModelOutput); // userInput here is for the *next* turn, but history is already primed
// Reset and set up a more realistic scenario for merging with existing history
chat = new GeminiChat(mockModelsModule, model, config, []);
chat = new GeminiChat(mockModelsModule, model, sessionId, config, []);
const firstUserInput: Content = {
role: 'user',
parts: [{ text: 'First user input' }],
@ -163,7 +164,7 @@ describe('GeminiChat', () => {
role: 'model',
parts: [{ text: 'Initial model answer.' }],
};
chat = new GeminiChat(mockModelsModule, model, config, [
chat = new GeminiChat(mockModelsModule, model, sessionId, config, [
initialUser,
initialModel,
]);

View File

@ -18,6 +18,7 @@ import {
import { retryWithBackoff } from '../utils/retry.js';
import { isFunctionResponse } from '../utils/messageInspectors.js';
import { ContentGenerator } from './contentGenerator.js';
import { Logger } from './logger.js';
/**
* Returns true if the response is valid, false otherwise.
@ -117,14 +118,17 @@ export class GeminiChat {
// A promise to represent the current state of the message being sent to the
// model.
private sendPromise: Promise<void> = Promise.resolve();
private logger: Logger;
constructor(
private readonly contentGenerator: ContentGenerator,
private readonly model: string,
sessionId: string,
private readonly config: GenerateContentConfig = {},
private history: Content[] = [],
) {
validateHistory(history);
this.logger = new Logger(sessionId);
}
/**

View File

@ -78,15 +78,20 @@ async function readLogFile(): Promise<LogEntry[]> {
}
}
vi.mock('../utils/session.js', () => ({
sessionId: 'test-session-id',
}));
describe('Logger', () => {
let logger: Logger;
const testSessionId = 'test-session-id';
beforeEach(async () => {
vi.resetAllMocks();
vi.useFakeTimers();
vi.setSystemTime(new Date('2025-01-01T12:00:00.000Z'));
await cleanupLogFile();
logger = new Logger();
logger = new Logger(testSessionId);
await logger.initialize();
});
@ -111,7 +116,7 @@ describe('Logger', () => {
/* ignore */
}
const newLogger = new Logger();
const newLogger = new Logger(testSessionId);
await newLogger.initialize();
const dirExists = await fs
@ -130,9 +135,8 @@ describe('Logger', () => {
});
it('should load existing logs and set correct messageId for the current session', async () => {
const fixedTime = new Date('2025-01-01T10:00:00.000Z');
vi.setSystemTime(fixedTime);
const currentSessionId = Math.floor(fixedTime.getTime() / 1000);
const currentSessionId = 'session-123';
const anotherSessionId = 'session-456';
const existingLogs: LogEntry[] = [
{
sessionId: currentSessionId,
@ -142,7 +146,7 @@ describe('Logger', () => {
message: 'Msg1',
},
{
sessionId: currentSessionId - 100,
sessionId: anotherSessionId,
messageId: 5,
timestamp: new Date('2025-01-01T09:00:00.000Z').toISOString(),
type: MessageSenderType.USER,
@ -158,7 +162,7 @@ describe('Logger', () => {
];
await fs.mkdir(path.join(process.cwd(), GEMINI_DIR), { recursive: true });
await fs.writeFile(TEST_LOG_FILE_PATH, JSON.stringify(existingLogs));
const newLogger = new Logger();
const newLogger = new Logger(currentSessionId);
await newLogger.initialize();
expect(newLogger['messageId']).toBe(2);
expect(newLogger['logs']).toEqual(existingLogs);
@ -166,11 +170,9 @@ describe('Logger', () => {
});
it('should set messageId to 0 for a new session if log file exists but has no logs for current session', async () => {
const fixedTime = new Date('2025-01-01T14:00:00.000Z');
vi.setSystemTime(fixedTime);
const existingLogs: LogEntry[] = [
{
sessionId: Math.floor(fixedTime.getTime() / 1000) - 500,
sessionId: 'some-other-session',
messageId: 5,
timestamp: new Date().toISOString(),
type: MessageSenderType.USER,
@ -179,7 +181,7 @@ describe('Logger', () => {
];
await fs.mkdir(path.join(process.cwd(), GEMINI_DIR), { recursive: true });
await fs.writeFile(TEST_LOG_FILE_PATH, JSON.stringify(existingLogs));
const newLogger = new Logger();
const newLogger = new Logger('a-new-session');
await newLogger.initialize();
expect(newLogger['messageId']).toBe(0);
newLogger.close();
@ -203,7 +205,7 @@ describe('Logger', () => {
const consoleDebugSpy = vi
.spyOn(console, 'debug')
.mockImplementation(() => {});
const newLogger = new Logger();
const newLogger = new Logger(testSessionId);
await newLogger.initialize();
expect(consoleDebugSpy).toHaveBeenCalledWith(
expect.stringContaining('Invalid JSON in log file'),
@ -232,7 +234,7 @@ describe('Logger', () => {
const consoleDebugSpy = vi
.spyOn(console, 'debug')
.mockImplementation(() => {});
const newLogger = new Logger();
const newLogger = new Logger(testSessionId);
await newLogger.initialize();
expect(consoleDebugSpy).toHaveBeenCalledWith(
`Log file at ${TEST_LOG_FILE_PATH} is not a valid JSON array. Starting with empty logs.`,
@ -259,7 +261,7 @@ describe('Logger', () => {
const logsFromFile = await readLogFile();
expect(logsFromFile.length).toBe(1);
expect(logsFromFile[0]).toMatchObject({
sessionId: logger['sessionId'],
sessionId: testSessionId,
messageId: 0,
type: MessageSenderType.USER,
message: 'Hello, world!',
@ -283,7 +285,8 @@ describe('Logger', () => {
});
it('should handle logger not initialized', async () => {
const uninitializedLogger = new Logger();
const uninitializedLogger = new Logger(testSessionId);
uninitializedLogger.close(); // Ensure it's treated as uninitialized
const consoleDebugSpy = vi
.spyOn(console, 'debug')
.mockImplementation(() => {});
@ -296,15 +299,13 @@ describe('Logger', () => {
});
it('should simulate concurrent writes from different logger instances to the same file', async () => {
const logger1 = new Logger(); // logger1
vi.setSystemTime(new Date('2025-01-01T13:00:00.000Z'));
const concurrentSessionId = 'concurrent-session';
const logger1 = new Logger(concurrentSessionId);
await logger1.initialize();
const s1 = logger1['sessionId'];
const logger2 = new Logger(); // logger2, will share same session if time is same
vi.setSystemTime(new Date('2025-01-01T13:00:00.000Z'));
const logger2 = new Logger(concurrentSessionId);
await logger2.initialize();
expect(logger2['sessionId']).toEqual(s1);
expect(logger2['sessionId']).toEqual(logger1['sessionId']);
// Log from logger1
await logger1.logMessage(MessageSenderType.USER, 'L1M1'); // L1 internal msgId becomes 1, writes {s1, 0}
@ -361,43 +362,34 @@ describe('Logger', () => {
});
describe('getPreviousUserMessages', () => {
it('should retrieve user messages, sorted newest first by session, then timestamp, then messageId', async () => {
const loggerSort = new Logger();
vi.setSystemTime(new Date('2025-01-01T10:00:00.000Z'));
it('should retrieve all user messages from logs, sorted newest first', async () => {
// This test now verifies that messages from different sessions are included
// and sorted correctly by timestamp, as the session-based sorting was removed.
const loggerSort = new Logger('session-1');
await loggerSort.initialize();
const s1 = loggerSort['sessionId']!;
await loggerSort.logMessage(MessageSenderType.USER, 'S1M0_ts100000'); // msgId 0
vi.advanceTimersByTime(10);
await loggerSort.logMessage(MessageSenderType.USER, 'S1M1_ts100010'); // msgId 1
loggerSort.close(); // Close to ensure next initialize starts a new session if time changed
vi.setSystemTime(new Date('2025-01-01T11:00:00.000Z'));
await loggerSort.initialize(); // Re-initialize for a new session
const s2 = loggerSort['sessionId']!;
expect(s2).not.toEqual(s1);
await loggerSort.logMessage(MessageSenderType.USER, 'S2M0_ts110000'); // msgId 0 for s2
vi.advanceTimersByTime(10);
vi.advanceTimersByTime(1000);
await loggerSort.logMessage(MessageSenderType.USER, 'S1M1_ts101000'); // msgId 1
vi.advanceTimersByTime(1000);
await loggerSort.logMessage(MessageSenderType.USER, 'S2M0_ts102000'); // msgId 0 for s2
vi.advanceTimersByTime(1000);
await loggerSort.logMessage(
'model' as MessageSenderType,
'S2_Model_ts110010',
'S2_Model_ts103000',
);
vi.advanceTimersByTime(10);
await loggerSort.logMessage(MessageSenderType.USER, 'S2M1_ts110020'); // msgId 1 for s2
vi.advanceTimersByTime(1000);
await loggerSort.logMessage(MessageSenderType.USER, 'S2M1_ts104000'); // msgId 1 for s2
loggerSort.close();
// To test the sorting thoroughly, especially the session part, we'll read the file
// as if it was written by multiple sessions and then initialize a new logger to load them.
const combinedLogs = await readLogFile();
const finalLogger = new Logger();
// Manually set its internal logs to simulate loading from a file with mixed sessions
finalLogger['logs'] = combinedLogs;
finalLogger['initialized'] = true; // Mark as initialized to allow getPreviousUserMessages to run
// A new logger will load all previous logs regardless of session
const finalLogger = new Logger('final-session');
await finalLogger.initialize();
const messages = await finalLogger.getPreviousUserMessages();
expect(messages).toEqual([
'S2M1_ts110020',
'S2M0_ts110000',
'S1M1_ts100010',
'S2M1_ts104000',
'S2M0_ts102000',
'S1M1_ts101000',
'S1M0_ts100000',
]);
finalLogger.close();
@ -410,7 +402,8 @@ describe('Logger', () => {
});
it('should return empty array if logger not initialized', async () => {
const uninitializedLogger = new Logger();
const uninitializedLogger = new Logger(testSessionId);
uninitializedLogger.close();
const messages = await uninitializedLogger.getPreviousUserMessages();
expect(messages).toEqual([]);
uninitializedLogger.close();
@ -443,7 +436,8 @@ describe('Logger', () => {
});
it('should not throw if logger is not initialized', async () => {
const uninitializedLogger = new Logger();
const uninitializedLogger = new Logger(testSessionId);
uninitializedLogger.close();
const consoleErrorSpy = vi
.spyOn(console, 'error')
.mockImplementation(() => {});
@ -521,7 +515,8 @@ describe('Logger', () => {
});
it('should return an empty array if logger is not initialized', async () => {
const uninitializedLogger = new Logger();
const uninitializedLogger = new Logger(testSessionId);
uninitializedLogger.close();
const consoleErrorSpy = vi
.spyOn(console, 'error')
.mockImplementation(() => {});

View File

@ -17,7 +17,7 @@ export enum MessageSenderType {
}
export interface LogEntry {
sessionId: number;
sessionId: string;
messageId: number;
timestamp: string;
type: MessageSenderType;
@ -28,12 +28,14 @@ export class Logger {
private geminiDir: string | undefined;
private logFilePath: string | undefined;
private checkpointFilePath: string | undefined;
private sessionId: number | undefined;
private sessionId: string | undefined;
private messageId = 0; // Instance-specific counter for the next messageId
private initialized = false;
private logs: LogEntry[] = []; // In-memory cache, ideally reflects the last known state of the file
constructor() {}
constructor(sessionId: string) {
this.sessionId = sessionId;
}
private async _readLogFile(): Promise<LogEntry[]> {
if (!this.logFilePath) {
@ -51,7 +53,7 @@ export class Logger {
}
return parsedLogs.filter(
(entry) =>
typeof entry.sessionId === 'number' &&
typeof entry.sessionId === 'string' &&
typeof entry.messageId === 'number' &&
typeof entry.timestamp === 'string' &&
typeof entry.type === 'string' &&
@ -93,7 +95,6 @@ export class Logger {
if (this.initialized) {
return;
}
this.sessionId = Math.floor(Date.now() / 1000);
this.geminiDir = path.resolve(process.cwd(), GEMINI_DIR);
this.logFilePath = path.join(this.geminiDir, LOG_FILE_NAME);
this.checkpointFilePath = path.join(this.geminiDir, CHECKPOINT_FILE_NAME);
@ -195,11 +196,9 @@ export class Logger {
return this.logs
.filter((entry) => entry.type === MessageSenderType.USER)
.sort((a, b) => {
if (b.sessionId !== a.sessionId) return b.sessionId - a.sessionId;
const dateA = new Date(a.timestamp).getTime();
const dateB = new Date(b.timestamp).getTime();
if (dateB !== dateA) return dateB - dateA;
return b.messageId - a.messageId;
return dateB - dateA;
})
.map((entry) => entry.message);
}

View File

@ -51,3 +51,4 @@ export * from './tools/mcp-tool.js';
// Export telemetry functions
export * from './telemetry/index.js';
export { sessionId } from './utils/session.js';

View File

@ -4,10 +4,7 @@
* SPDX-License-Identifier: Apache-2.0
*/
import { randomUUID } from 'crypto';
export const SERVICE_NAME = 'gemini-cli';
export const sessionId = randomUUID();
export const EVENT_USER_PROMPT = 'gemini_code.user_prompt';
export const EVENT_TOOL_CALL = 'gemini_code.tool_call';

View File

@ -28,4 +28,3 @@ export {
} from './types.js';
export { SpanStatusCode, ValueType } from '@opentelemetry/api';
export { SemanticAttributes } from '@opentelemetry/semantic-conventions';
export { sessionId } from './constants.js';

View File

@ -26,7 +26,7 @@ import {
} from '@opentelemetry/sdk-metrics';
import { HttpInstrumentation } from '@opentelemetry/instrumentation-http';
import { Config } from '../config/config.js';
import { SERVICE_NAME, sessionId } from './constants.js';
import { SERVICE_NAME } from './constants.js';
import { initializeMetrics } from './metrics.js';
import { logCliConfiguration } from './loggers.js';
@ -68,7 +68,7 @@ export function initializeTelemetry(config: Config): void {
const resource = new Resource({
[SemanticResourceAttributes.SERVICE_NAME]: SERVICE_NAME,
[SemanticResourceAttributes.SERVICE_VERSION]: process.version,
'session.id': sessionId,
'session.id': config.getSessionId(),
});
const otlpEndpoint = config.getTelemetryOtlpEndpoint();

View File

@ -136,6 +136,7 @@ const baseConfigParams: ConfigParameters = {
userMemory: '',
geminiMdFileCount: 0,
approvalMode: ApprovalMode.DEFAULT,
sessionId: 'test-session-id',
};
describe('ToolRegistry', () => {

View File

@ -71,6 +71,7 @@ describe('checkNextSpeaker', () => {
chatInstance = new GeminiChat(
mockModelsInstance, // This is the instance returned by mockGoogleGenAIInstance.getGenerativeModel
'gemini-pro', // model name
'test-session-id',
{},
[], // initial history
);

View File

@ -0,0 +1,9 @@
/**
* @license
* Copyright 2025 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { randomUUID } from 'crypto';
export const sessionId = randomUUID();