diff options
Diffstat (limited to 'packages/server/src')
| -rw-r--r-- | packages/server/src/core/logger.test.ts | 197 | ||||
| -rw-r--r-- | packages/server/src/core/logger.ts | 131 | ||||
| -rw-r--r-- | packages/server/src/index.ts | 1 |
3 files changed, 329 insertions, 0 deletions
diff --git a/packages/server/src/core/logger.test.ts b/packages/server/src/core/logger.test.ts new file mode 100644 index 00000000..9b4f4555 --- /dev/null +++ b/packages/server/src/core/logger.test.ts @@ -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(); + }); + }); +}); diff --git a/packages/server/src/core/logger.ts b/packages/server/src/core/logger.ts new file mode 100644 index 00000000..d12d4240 --- /dev/null +++ b/packages/server/src/core/logger.ts @@ -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; + } + } +} diff --git a/packages/server/src/index.ts b/packages/server/src/index.ts index c90ef2fd..788fe6e4 100644 --- a/packages/server/src/index.ts +++ b/packages/server/src/index.ts @@ -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'; |
