mirror of
https://gh-proxy.org/https://github.com/klinsmaya/LobsterAI.git
synced 2026-11-10 18:26:44 +08:00
444 lines
14 KiB
TypeScript
444 lines
14 KiB
TypeScript
import { app } from 'electron';
|
|
import { EventEmitter } from 'events';
|
|
import crypto from 'crypto';
|
|
import fs from 'fs';
|
|
import path from 'path';
|
|
import initSqlJs, { Database, SqlJsStatic } from 'sql.js';
|
|
import { DB_FILENAME } from './appConstants';
|
|
|
|
type ChangePayload<T = unknown> = {
|
|
key: string;
|
|
newValue: T | undefined;
|
|
oldValue: T | undefined;
|
|
};
|
|
|
|
const USER_MEMORIES_MIGRATION_KEY = 'userMemories.migration.v1.completed';
|
|
|
|
// Pre-read the sql.js WASM binary from disk.
|
|
// Using fs.readFileSync (which handles non-ASCII paths via Windows wide-char APIs)
|
|
// and passing the buffer directly to initSqlJs bypasses Emscripten's file loading,
|
|
// which can fail or hang when the install path contains Chinese characters on Windows.
|
|
function loadWasmBinary(): ArrayBuffer {
|
|
const wasmPath = app.isPackaged
|
|
? path.join(
|
|
process.resourcesPath,
|
|
'app.asar.unpacked/node_modules/sql.js/dist/sql-wasm.wasm'
|
|
)
|
|
: path.join(app.getAppPath(), 'node_modules/sql.js/dist/sql-wasm.wasm');
|
|
const buf = fs.readFileSync(wasmPath);
|
|
return buf.buffer.slice(buf.byteOffset, buf.byteOffset + buf.byteLength);
|
|
}
|
|
|
|
export class SqliteStore {
|
|
private db: Database;
|
|
private dbPath: string;
|
|
private emitter = new EventEmitter();
|
|
private static sqlPromise: Promise<SqlJsStatic> | null = null;
|
|
|
|
private constructor(db: Database, dbPath: string) {
|
|
this.db = db;
|
|
this.dbPath = dbPath;
|
|
}
|
|
|
|
static async create(userDataPath?: string): Promise<SqliteStore> {
|
|
const basePath = userDataPath ?? app.getPath('userData');
|
|
const dbPath = path.join(basePath, DB_FILENAME);
|
|
|
|
// Initialize SQL.js with WASM file path (cached promise for reuse)
|
|
if (!SqliteStore.sqlPromise) {
|
|
const wasmBinary = loadWasmBinary();
|
|
SqliteStore.sqlPromise = initSqlJs({
|
|
wasmBinary,
|
|
});
|
|
}
|
|
const SQL = await SqliteStore.sqlPromise;
|
|
|
|
// Load existing database or create new one
|
|
let db: Database;
|
|
if (fs.existsSync(dbPath)) {
|
|
const buffer = fs.readFileSync(dbPath);
|
|
db = new SQL.Database(buffer);
|
|
} else {
|
|
db = new SQL.Database();
|
|
}
|
|
|
|
const store = new SqliteStore(db, dbPath);
|
|
store.initializeTables(basePath);
|
|
return store;
|
|
}
|
|
|
|
private initializeTables(basePath: string) {
|
|
this.db.run(`
|
|
CREATE TABLE IF NOT EXISTS kv (
|
|
key TEXT PRIMARY KEY,
|
|
value TEXT NOT NULL,
|
|
updated_at INTEGER NOT NULL
|
|
);
|
|
`);
|
|
|
|
// Create cowork tables
|
|
this.db.run(`
|
|
CREATE TABLE IF NOT EXISTS cowork_sessions (
|
|
id TEXT PRIMARY KEY,
|
|
title TEXT NOT NULL,
|
|
claude_session_id TEXT,
|
|
status TEXT NOT NULL DEFAULT 'idle',
|
|
pinned INTEGER NOT NULL DEFAULT 0,
|
|
cwd TEXT NOT NULL,
|
|
system_prompt TEXT NOT NULL DEFAULT '',
|
|
execution_mode TEXT,
|
|
created_at INTEGER NOT NULL,
|
|
updated_at INTEGER NOT NULL
|
|
);
|
|
`);
|
|
|
|
this.db.run(`
|
|
CREATE TABLE IF NOT EXISTS cowork_messages (
|
|
id TEXT PRIMARY KEY,
|
|
session_id TEXT NOT NULL,
|
|
type TEXT NOT NULL,
|
|
content TEXT NOT NULL,
|
|
metadata TEXT,
|
|
created_at INTEGER NOT NULL,
|
|
sequence INTEGER,
|
|
FOREIGN KEY (session_id) REFERENCES cowork_sessions(id) ON DELETE CASCADE
|
|
);
|
|
`);
|
|
|
|
this.db.run(`
|
|
CREATE INDEX IF NOT EXISTS idx_cowork_messages_session_id ON cowork_messages(session_id);
|
|
`);
|
|
|
|
this.db.run(`
|
|
CREATE TABLE IF NOT EXISTS cowork_config (
|
|
key TEXT PRIMARY KEY,
|
|
value TEXT NOT NULL,
|
|
updated_at INTEGER NOT NULL
|
|
);
|
|
`);
|
|
|
|
this.db.run(`
|
|
CREATE TABLE IF NOT EXISTS user_memories (
|
|
id TEXT PRIMARY KEY,
|
|
text TEXT NOT NULL,
|
|
fingerprint TEXT NOT NULL,
|
|
confidence REAL NOT NULL DEFAULT 0.75,
|
|
is_explicit INTEGER NOT NULL DEFAULT 0,
|
|
status TEXT NOT NULL DEFAULT 'created',
|
|
created_at INTEGER NOT NULL,
|
|
updated_at INTEGER NOT NULL,
|
|
last_used_at INTEGER
|
|
);
|
|
`);
|
|
|
|
this.db.run(`
|
|
CREATE TABLE IF NOT EXISTS user_memory_sources (
|
|
id TEXT PRIMARY KEY,
|
|
memory_id TEXT NOT NULL,
|
|
session_id TEXT,
|
|
message_id TEXT,
|
|
role TEXT NOT NULL DEFAULT 'system',
|
|
is_active INTEGER NOT NULL DEFAULT 1,
|
|
created_at INTEGER NOT NULL,
|
|
FOREIGN KEY (memory_id) REFERENCES user_memories(id) ON DELETE CASCADE
|
|
);
|
|
`);
|
|
|
|
this.db.run(`
|
|
CREATE INDEX IF NOT EXISTS idx_user_memories_status_updated_at
|
|
ON user_memories(status, updated_at DESC);
|
|
`);
|
|
this.db.run(`
|
|
CREATE INDEX IF NOT EXISTS idx_user_memories_fingerprint
|
|
ON user_memories(fingerprint);
|
|
`);
|
|
this.db.run(`
|
|
CREATE INDEX IF NOT EXISTS idx_user_memory_sources_session_id
|
|
ON user_memory_sources(session_id, is_active);
|
|
`);
|
|
this.db.run(`
|
|
CREATE INDEX IF NOT EXISTS idx_user_memory_sources_memory_id
|
|
ON user_memory_sources(memory_id, is_active);
|
|
`);
|
|
|
|
// Create MCP servers table
|
|
this.db.run(`
|
|
CREATE TABLE IF NOT EXISTS mcp_servers (
|
|
id TEXT PRIMARY KEY,
|
|
name TEXT NOT NULL UNIQUE,
|
|
description TEXT NOT NULL DEFAULT '',
|
|
enabled INTEGER NOT NULL DEFAULT 1,
|
|
transport_type TEXT NOT NULL DEFAULT 'stdio',
|
|
config_json TEXT NOT NULL DEFAULT '{}',
|
|
created_at INTEGER NOT NULL,
|
|
updated_at INTEGER NOT NULL
|
|
);
|
|
`);
|
|
|
|
// Migrations - safely add columns if they don't exist
|
|
try {
|
|
// Check if execution_mode column exists
|
|
const colsResult = this.db.exec("PRAGMA table_info(cowork_sessions);");
|
|
const columns = colsResult[0]?.values.map((row) => row[1]) || [];
|
|
|
|
if (!columns.includes('execution_mode')) {
|
|
this.db.run('ALTER TABLE cowork_sessions ADD COLUMN execution_mode TEXT;');
|
|
this.save();
|
|
}
|
|
|
|
if (!columns.includes('pinned')) {
|
|
this.db.run('ALTER TABLE cowork_sessions ADD COLUMN pinned INTEGER NOT NULL DEFAULT 0;');
|
|
this.save();
|
|
}
|
|
|
|
if (!columns.includes('active_skill_ids')) {
|
|
this.db.run('ALTER TABLE cowork_sessions ADD COLUMN active_skill_ids TEXT;');
|
|
this.save();
|
|
}
|
|
|
|
// Migration: Add sequence column to cowork_messages
|
|
const msgColsResult = this.db.exec("PRAGMA table_info(cowork_messages);");
|
|
const msgColumns = msgColsResult[0]?.values.map((row) => row[1]) || [];
|
|
|
|
if (!msgColumns.includes('sequence')) {
|
|
this.db.run('ALTER TABLE cowork_messages ADD COLUMN sequence INTEGER');
|
|
|
|
// 为现有消息按 created_at 和 ROWID 分配序列号
|
|
this.db.run(`
|
|
WITH numbered AS (
|
|
SELECT id, ROW_NUMBER() OVER (
|
|
PARTITION BY session_id
|
|
ORDER BY created_at ASC, ROWID ASC
|
|
) as seq
|
|
FROM cowork_messages
|
|
)
|
|
UPDATE cowork_messages
|
|
SET sequence = (SELECT seq FROM numbered WHERE numbered.id = cowork_messages.id)
|
|
`);
|
|
|
|
this.save();
|
|
}
|
|
} catch {
|
|
// Column already exists or migration not needed.
|
|
}
|
|
|
|
try {
|
|
this.db.run('UPDATE cowork_sessions SET pinned = 0 WHERE pinned IS NULL;');
|
|
} catch {
|
|
// Column might not exist yet.
|
|
}
|
|
|
|
try {
|
|
this.db.run(`UPDATE cowork_sessions SET execution_mode = 'local' WHERE execution_mode = 'container';`);
|
|
this.db.run(`
|
|
UPDATE cowork_config
|
|
SET value = 'local'
|
|
WHERE key = 'executionMode' AND value = 'container';
|
|
`);
|
|
} catch (error) {
|
|
console.warn('Failed to migrate cowork execution mode:', error);
|
|
}
|
|
|
|
this.migrateLegacyMemoryFileToUserMemories();
|
|
this.migrateFromElectronStore(basePath);
|
|
this.save();
|
|
}
|
|
|
|
save() {
|
|
const data = this.db.export();
|
|
const buffer = Buffer.from(data);
|
|
fs.writeFileSync(this.dbPath, buffer);
|
|
}
|
|
|
|
onDidChange<T = unknown>(key: string, callback: (newValue: T | undefined, oldValue: T | undefined) => void) {
|
|
const handler = (payload: ChangePayload<T>) => {
|
|
if (payload.key !== key) return;
|
|
callback(payload.newValue, payload.oldValue);
|
|
};
|
|
this.emitter.on('change', handler);
|
|
return () => this.emitter.off('change', handler);
|
|
}
|
|
|
|
get<T = unknown>(key: string): T | undefined {
|
|
const result = this.db.exec('SELECT value FROM kv WHERE key = ?', [key]);
|
|
if (!result[0]?.values[0]) return undefined;
|
|
const value = result[0].values[0][0] as string;
|
|
try {
|
|
return JSON.parse(value) as T;
|
|
} catch (error) {
|
|
console.warn(`Failed to parse store value for ${key}`, error);
|
|
return undefined;
|
|
}
|
|
}
|
|
|
|
set<T = unknown>(key: string, value: T): void {
|
|
const oldValue = this.get<T>(key);
|
|
const now = Date.now();
|
|
this.db.run(`
|
|
INSERT INTO kv (key, value, updated_at)
|
|
VALUES (?, ?, ?)
|
|
ON CONFLICT(key) DO UPDATE SET
|
|
value = excluded.value,
|
|
updated_at = excluded.updated_at
|
|
`, [key, JSON.stringify(value), now]);
|
|
this.save();
|
|
this.emitter.emit('change', { key, newValue: value, oldValue } as ChangePayload<T>);
|
|
}
|
|
|
|
delete(key: string): void {
|
|
const oldValue = this.get(key);
|
|
this.db.run('DELETE FROM kv WHERE key = ?', [key]);
|
|
this.save();
|
|
this.emitter.emit('change', { key, newValue: undefined, oldValue } as ChangePayload);
|
|
}
|
|
|
|
// Expose database for cowork operations
|
|
getDatabase(): Database {
|
|
return this.db;
|
|
}
|
|
|
|
// Expose save method for external use (e.g., CoworkStore)
|
|
getSaveFunction(): () => void {
|
|
return () => this.save();
|
|
}
|
|
|
|
private tryReadLegacyMemoryText(): string {
|
|
const candidates = [
|
|
path.join(process.cwd(), 'MEMORY.md'),
|
|
path.join(app.getAppPath(), 'MEMORY.md'),
|
|
path.join(process.cwd(), 'memory.md'),
|
|
path.join(app.getAppPath(), 'memory.md'),
|
|
];
|
|
|
|
for (const candidate of candidates) {
|
|
try {
|
|
if (fs.existsSync(candidate) && fs.statSync(candidate).isFile()) {
|
|
return fs.readFileSync(candidate, 'utf8');
|
|
}
|
|
} catch {
|
|
// Skip unreadable candidates.
|
|
}
|
|
}
|
|
return '';
|
|
}
|
|
|
|
private parseLegacyMemoryEntries(raw: string): string[] {
|
|
const normalized = raw.replace(/```[\s\S]*?```/g, ' ');
|
|
const lines = normalized.split(/\r?\n/);
|
|
const entries: string[] = [];
|
|
const seen = new Set<string>();
|
|
|
|
for (const line of lines) {
|
|
const match = line.trim().match(/^-+\s*(?:\[[^\]]+\]\s*)?(.+)$/);
|
|
if (!match?.[1]) continue;
|
|
const text = match[1].replace(/\s+/g, ' ').trim();
|
|
if (!text || text.length < 6) continue;
|
|
if (/^\(empty\)$/i.test(text)) continue;
|
|
const key = text.toLowerCase();
|
|
if (seen.has(key)) continue;
|
|
seen.add(key);
|
|
entries.push(text.length > 360 ? `${text.slice(0, 359)}…` : text);
|
|
}
|
|
|
|
return entries.slice(0, 200);
|
|
}
|
|
|
|
private memoryFingerprint(text: string): string {
|
|
const normalized = text
|
|
.toLowerCase()
|
|
.replace(/[^\p{L}\p{N}\s]/gu, ' ')
|
|
.replace(/\s+/g, ' ')
|
|
.trim();
|
|
return crypto.createHash('sha1').update(normalized).digest('hex');
|
|
}
|
|
|
|
private migrateLegacyMemoryFileToUserMemories(): void {
|
|
if (this.get<string>(USER_MEMORIES_MIGRATION_KEY) === '1') {
|
|
return;
|
|
}
|
|
|
|
const content = this.tryReadLegacyMemoryText();
|
|
if (!content.trim()) {
|
|
this.set(USER_MEMORIES_MIGRATION_KEY, '1');
|
|
return;
|
|
}
|
|
|
|
const entries = this.parseLegacyMemoryEntries(content);
|
|
if (entries.length === 0) {
|
|
this.set(USER_MEMORIES_MIGRATION_KEY, '1');
|
|
return;
|
|
}
|
|
|
|
const now = Date.now();
|
|
this.db.run('BEGIN TRANSACTION;');
|
|
try {
|
|
for (const text of entries) {
|
|
const fingerprint = this.memoryFingerprint(text);
|
|
const existing = this.db.exec(
|
|
`SELECT id FROM user_memories WHERE fingerprint = ? AND status != 'deleted' LIMIT 1`,
|
|
[fingerprint]
|
|
);
|
|
if (existing[0]?.values?.[0]?.[0]) {
|
|
continue;
|
|
}
|
|
|
|
const memoryId = crypto.randomUUID();
|
|
this.db.run(`
|
|
INSERT INTO user_memories (
|
|
id, text, fingerprint, confidence, is_explicit, status, created_at, updated_at, last_used_at
|
|
) VALUES (?, ?, ?, ?, 1, 'created', ?, ?, NULL)
|
|
`, [memoryId, text, fingerprint, 0.9, now, now]);
|
|
|
|
this.db.run(`
|
|
INSERT INTO user_memory_sources (id, memory_id, session_id, message_id, role, is_active, created_at)
|
|
VALUES (?, ?, NULL, NULL, 'system', 1, ?)
|
|
`, [crypto.randomUUID(), memoryId, now]);
|
|
}
|
|
|
|
this.db.run('COMMIT;');
|
|
} catch (error) {
|
|
this.db.run('ROLLBACK;');
|
|
console.warn('Failed to migrate legacy MEMORY.md entries:', error);
|
|
}
|
|
|
|
this.set(USER_MEMORIES_MIGRATION_KEY, '1');
|
|
}
|
|
|
|
private migrateFromElectronStore(userDataPath: string) {
|
|
const result = this.db.exec('SELECT COUNT(*) as count FROM kv');
|
|
const count = result[0]?.values[0]?.[0] as number;
|
|
if (count > 0) return;
|
|
|
|
const legacyPath = path.join(userDataPath, 'config.json');
|
|
if (!fs.existsSync(legacyPath)) return;
|
|
|
|
try {
|
|
const raw = fs.readFileSync(legacyPath, 'utf8');
|
|
const data = JSON.parse(raw) as Record<string, unknown>;
|
|
if (!data || typeof data !== 'object') return;
|
|
|
|
const entries = Object.entries(data);
|
|
if (!entries.length) return;
|
|
|
|
const now = Date.now();
|
|
this.db.run('BEGIN TRANSACTION;');
|
|
try {
|
|
entries.forEach(([key, value]) => {
|
|
this.db.run(`
|
|
INSERT INTO kv (key, value, updated_at)
|
|
VALUES (?, ?, ?)
|
|
`, [key, JSON.stringify(value), now]);
|
|
});
|
|
this.db.run('COMMIT;');
|
|
this.save();
|
|
console.info(`Migrated ${entries.length} entries from electron-store.`);
|
|
} catch (error) {
|
|
this.db.run('ROLLBACK;');
|
|
throw error;
|
|
}
|
|
} catch (error) {
|
|
console.warn('Failed to migrate electron-store data:', error);
|
|
}
|
|
}
|
|
}
|