mirror of
https://github.com/QwenLM/qwen-code.git
synced 2025-12-20 16:57:46 +00:00
Add Logger for command history (#435)
This commit is contained in:
197
packages/server/src/core/logger.test.ts
Normal file
197
packages/server/src/core/logger.test.ts
Normal file
@@ -0,0 +1,197 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
|
||||
import { Logger, MessageSenderType } from './logger.js';
|
||||
|
||||
// Mocks
|
||||
const mockDb = {
|
||||
exec: vi.fn((_sql, callback) => callback?.(null)),
|
||||
all: vi.fn((_sql, _params, callback) => callback?.(null, [])),
|
||||
run: vi.fn((_sql, _params, callback) => callback?.(null)),
|
||||
close: vi.fn((callback) => callback?.(null)),
|
||||
};
|
||||
|
||||
vi.mock('sqlite3', () => ({
|
||||
Database: vi.fn((_dbPath, _options, callback) => {
|
||||
process.nextTick(() => callback?.(null));
|
||||
return mockDb;
|
||||
}),
|
||||
default: {
|
||||
Database: vi.fn((_dbPath, _options, callback) => {
|
||||
process.nextTick(() => callback?.(null));
|
||||
return mockDb;
|
||||
}),
|
||||
},
|
||||
}));
|
||||
|
||||
describe('Logger', () => {
|
||||
let logger: Logger;
|
||||
|
||||
beforeEach(async () => {
|
||||
vi.resetAllMocks();
|
||||
|
||||
// Get a new instance for each test to ensure isolation,
|
||||
logger = new Logger();
|
||||
// We need to wait for the async initialize to complete
|
||||
await logger.initialize().catch((err) => {
|
||||
console.error('Error initializing logger:', err);
|
||||
});
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
vi.restoreAllMocks();
|
||||
logger.close(); // Close the database connection after each test
|
||||
});
|
||||
|
||||
describe('initialize', () => {
|
||||
it('should execute create tables if not exists', async () => {
|
||||
expect(mockDb.exec).toHaveBeenCalledWith(
|
||||
expect.stringMatching(/CREATE TABLE IF NOT EXISTS messages/),
|
||||
expect.any(Function),
|
||||
);
|
||||
});
|
||||
|
||||
it('should be idempotent', async () => {
|
||||
mockDb.exec.mockClear();
|
||||
|
||||
await logger.initialize(); // Second call
|
||||
|
||||
expect(mockDb.exec).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
||||
describe('logMessage', () => {
|
||||
it('should insert a message into the database', async () => {
|
||||
const type = MessageSenderType.USER;
|
||||
const message = 'Hello, world!';
|
||||
await logger.logMessage(type, message);
|
||||
expect(mockDb.run).toHaveBeenCalledWith(
|
||||
"INSERT INTO messages (session_id, message_id, type, message, timestamp) VALUES (?, ?, ?, ?, datetime('now'))",
|
||||
[expect.any(Number), 0, type, message], // sessionId, messageId, type, message
|
||||
expect.any(Function),
|
||||
);
|
||||
});
|
||||
|
||||
it('should increment messageId for subsequent messages', async () => {
|
||||
await logger.logMessage(MessageSenderType.USER, 'First message');
|
||||
expect(mockDb.run).toHaveBeenCalledWith(
|
||||
expect.any(String),
|
||||
[expect.any(Number), 0, MessageSenderType.USER, 'First message'],
|
||||
expect.any(Function),
|
||||
);
|
||||
await logger.logMessage(MessageSenderType.USER, 'Second message');
|
||||
expect(mockDb.run).toHaveBeenCalledWith(
|
||||
expect.any(String),
|
||||
[expect.any(Number), 1, MessageSenderType.USER, 'Second message'], // messageId is now 1
|
||||
expect.any(Function),
|
||||
);
|
||||
});
|
||||
|
||||
it('should handle database not initialized', async () => {
|
||||
const uninitializedLogger = new Logger();
|
||||
// uninitializedLogger.initialize() is not called
|
||||
const consoleErrorSpy = vi
|
||||
.spyOn(console, 'error')
|
||||
.mockImplementation(() => {});
|
||||
|
||||
await uninitializedLogger.logMessage(MessageSenderType.USER, 'test');
|
||||
|
||||
expect(consoleErrorSpy).toHaveBeenCalledWith('Database not initialized.');
|
||||
expect(mockDb.run).not.toHaveBeenCalled();
|
||||
consoleErrorSpy.mockRestore();
|
||||
});
|
||||
|
||||
it('should handle error during db.run', async () => {
|
||||
const error = new Error('db.run failed');
|
||||
mockDb.run.mockImplementationOnce(
|
||||
// eslint-disable-next-line @typescript-eslint/no-explicit-any
|
||||
(_sql: any, _params: any, callback: any) => callback?.(error),
|
||||
);
|
||||
|
||||
await expect(
|
||||
logger.logMessage(MessageSenderType.USER, 'test'),
|
||||
).rejects.toThrow('db.run failed');
|
||||
});
|
||||
});
|
||||
|
||||
describe('getPreviousUserMessages', () => {
|
||||
it('should query the database for messages', async () => {
|
||||
mockDb.all.mockImplementationOnce(
|
||||
// eslint-disable-next-line @typescript-eslint/no-explicit-any
|
||||
(_sql: any, params: any, callback: any) =>
|
||||
callback?.(null, [{ message: 'msg1' }, { message: 'msg2' }]),
|
||||
);
|
||||
|
||||
const messages = await logger.getPreviousUserMessages();
|
||||
|
||||
expect(mockDb.all).toHaveBeenCalledWith(
|
||||
expect.stringMatching(/SELECT message FROM messages/),
|
||||
[],
|
||||
expect.any(Function),
|
||||
);
|
||||
expect(messages).toEqual(['msg1', 'msg2']);
|
||||
});
|
||||
|
||||
it('should handle database not initialized', async () => {
|
||||
const uninitializedLogger = new Logger();
|
||||
// uninitializedLogger.initialize() is not called
|
||||
const consoleErrorSpy = vi
|
||||
.spyOn(console, 'error')
|
||||
.mockImplementation(() => {});
|
||||
|
||||
const messages = await uninitializedLogger.getPreviousUserMessages();
|
||||
|
||||
expect(consoleErrorSpy).toHaveBeenCalledWith('Database not initialized.');
|
||||
expect(messages).toEqual([]);
|
||||
expect(mockDb.all).not.toHaveBeenCalled();
|
||||
consoleErrorSpy.mockRestore();
|
||||
});
|
||||
|
||||
it('should handle error during db.all', async () => {
|
||||
const error = new Error('db.all failed');
|
||||
mockDb.all.mockImplementationOnce(
|
||||
// eslint-disable-next-line @typescript-eslint/no-explicit-any
|
||||
(_sql: any, _params: any, callback: any) => callback?.(error, []),
|
||||
);
|
||||
|
||||
await expect(logger.getPreviousUserMessages()).rejects.toThrow(
|
||||
'db.all failed',
|
||||
);
|
||||
});
|
||||
});
|
||||
|
||||
describe('close', () => {
|
||||
it('should close the database connection', () => {
|
||||
logger.close();
|
||||
expect(mockDb.close).toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('should handle database not initialized', () => {
|
||||
const uninitializedLogger = new Logger();
|
||||
// uninitializedLogger.initialize() is not called
|
||||
uninitializedLogger.close();
|
||||
expect(() => uninitializedLogger.close()).not.toThrow();
|
||||
});
|
||||
|
||||
it('should handle error during db.close', () => {
|
||||
const error = new Error('db.close failed');
|
||||
mockDb.close.mockImplementationOnce((callback: (error: Error) => void) =>
|
||||
callback?.(error),
|
||||
);
|
||||
const consoleErrorSpy = vi
|
||||
.spyOn(console, 'error')
|
||||
.mockImplementation(() => {});
|
||||
|
||||
logger.close();
|
||||
expect(consoleErrorSpy).toHaveBeenCalledWith(
|
||||
'Error closing database:',
|
||||
error.message,
|
||||
);
|
||||
consoleErrorSpy.mockRestore();
|
||||
});
|
||||
});
|
||||
});
|
||||
131
packages/server/src/core/logger.ts
Normal file
131
packages/server/src/core/logger.ts
Normal file
@@ -0,0 +1,131 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2025 Google LLC
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*/
|
||||
|
||||
import path from 'node:path';
|
||||
import sqlite3 from 'sqlite3';
|
||||
import { promises as fs } from 'node:fs';
|
||||
|
||||
const GEMINI_DIR = '.gemini';
|
||||
const DB_NAME = 'logs.db';
|
||||
const CREATE_TABLE_SQL = `
|
||||
CREATE TABLE IF NOT EXISTS messages (
|
||||
session_id INTEGER,
|
||||
message_id INTEGER,
|
||||
timestamp TEXT,
|
||||
type TEXT,
|
||||
message TEXT
|
||||
);`;
|
||||
|
||||
export enum MessageSenderType {
|
||||
USER = 'user',
|
||||
}
|
||||
|
||||
export class Logger {
|
||||
private db: sqlite3.Database | undefined;
|
||||
private sessionId: number | undefined;
|
||||
private messageId: number | undefined;
|
||||
|
||||
constructor() {}
|
||||
|
||||
async initialize(): Promise<void> {
|
||||
if (this.db) {
|
||||
return;
|
||||
}
|
||||
|
||||
this.sessionId = Math.floor(Date.now() / 1000);
|
||||
this.messageId = 0;
|
||||
|
||||
// Could be cleaner if our sqlite package supported promises.
|
||||
return new Promise((resolve, reject) => {
|
||||
const DB_DIR = path.resolve(process.cwd(), GEMINI_DIR);
|
||||
const DB_PATH = path.join(DB_DIR, DB_NAME);
|
||||
fs.mkdir(DB_DIR, { recursive: true })
|
||||
.then(() => {
|
||||
this.db = new sqlite3.Database(
|
||||
DB_PATH,
|
||||
sqlite3.OPEN_READWRITE |
|
||||
sqlite3.OPEN_CREATE |
|
||||
sqlite3.OPEN_FULLMUTEX,
|
||||
(err: Error | null) => {
|
||||
if (err) {
|
||||
reject(err);
|
||||
}
|
||||
|
||||
// Read and execute the SQL script in create_tables.sql
|
||||
this.db?.exec(CREATE_TABLE_SQL, (err: Error | null) => {
|
||||
if (err) {
|
||||
this.db?.close();
|
||||
reject(err);
|
||||
}
|
||||
resolve();
|
||||
});
|
||||
},
|
||||
);
|
||||
})
|
||||
.catch(reject);
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Get list of previous user inputs sorted most recent first.
|
||||
* @returns list of messages.
|
||||
*/
|
||||
async getPreviousUserMessages(): Promise<string[]> {
|
||||
if (!this.db) {
|
||||
console.error('Database not initialized.');
|
||||
return [];
|
||||
}
|
||||
|
||||
return new Promise((resolve, reject) => {
|
||||
// Most recent messages first
|
||||
const query = `SELECT message FROM messages
|
||||
WHERE type = '${MessageSenderType.USER}'
|
||||
ORDER BY session_id DESC, message_id DESC`;
|
||||
// eslint-disable-next-line @typescript-eslint/no-explicit-any
|
||||
this.db!.all(query, [], (err: Error | null, rows: any[]) => {
|
||||
if (err) {
|
||||
reject(err);
|
||||
} else {
|
||||
resolve(rows.map((row) => row.message));
|
||||
}
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
async logMessage(type: MessageSenderType, message: string): Promise<void> {
|
||||
if (!this.db) {
|
||||
console.error('Database not initialized.');
|
||||
return;
|
||||
}
|
||||
|
||||
return new Promise((resolve, reject) => {
|
||||
const query = `INSERT INTO messages (session_id, message_id, type, message, timestamp) VALUES (?, ?, ?, ?, datetime('now'))`;
|
||||
this.messageId = this.messageId! + 1;
|
||||
this.db!.run(
|
||||
query,
|
||||
[this.sessionId || 0, this.messageId - 1, type, message],
|
||||
(err: Error | null) => {
|
||||
if (err) {
|
||||
reject(err);
|
||||
} else {
|
||||
resolve();
|
||||
}
|
||||
},
|
||||
);
|
||||
});
|
||||
}
|
||||
|
||||
close(): void {
|
||||
if (this.db) {
|
||||
this.db.close((err: Error | null) => {
|
||||
if (err) {
|
||||
console.error('Error closing database:', err.message);
|
||||
}
|
||||
});
|
||||
this.db = undefined;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -9,6 +9,7 @@ export * from './config/config.js';
|
||||
|
||||
// Export Core Logic
|
||||
export * from './core/client.js';
|
||||
export * from './core/logger.js';
|
||||
export * from './core/prompts.js';
|
||||
export * from './core/turn.js';
|
||||
export * from './core/geminiRequest.js';
|
||||
|
||||
Reference in New Issue
Block a user