📄 database.ts  •  4341 bytes
/**
 * CmdCode 向量记忆系统 - 数据库初始化(使用 Bun 内置 SQLite)
 */
import { t } from '../i18n'
import { Database } from 'bun:sqlite'
import { join } from 'path'
import { existsSync, mkdirSync } from 'fs'
import os from 'os'

const DB_PATH = join(os.homedir(), '.cmdcode', 'memory', 'cmdcode.db')

// 确保目录存在
const dir = join(os.homedir(), '.cmdcode', 'memory')
if (!existsSync(dir)) {
  mkdirSync(dir, { recursive: true })
}

// 数据库实例
let db: Database | null = null
let isClosed = false  // P3 #32: 追踪关闭状态

/** 获取数据库实例(单例) - P3 #32: 支持关闭后重开 */
export function getDb(): Database {
  if (isClosed && !db) {
    // 关闭后首次调用,重置状态并重开
    isClosed = false
  }
  if (!db) {
    db = new Database(DB_PATH)
    db.run('PRAGMA foreign_keys = ON')
    initializeDatabase(db)
  }
  return db
}

/** 初始化数据库表 */
export function initializeDatabase(database?: Database): void {
  const dbInstance = database || getDb()

  dbInstance.run(`
    -- 会话表
    CREATE TABLE IF NOT EXISTS sessions (
      id TEXT PRIMARY KEY,
      title TEXT,
      created_at TEXT DEFAULT (datetime('now')),
      updated_at TEXT DEFAULT (datetime('now'))
    )
  `)

  dbInstance.run(`
    -- 消息表
    CREATE TABLE IF NOT EXISTS messages (
      id INTEGER PRIMARY KEY AUTOINCREMENT,
      session_id TEXT NOT NULL REFERENCES sessions(id) ON DELETE CASCADE,
      role TEXT NOT NULL CHECK(role IN ('user', 'assistant', 'system', 'tool')),
      content TEXT NOT NULL,
      created_at TEXT DEFAULT (datetime('now'))
    )
  `)

  dbInstance.run(`
    -- FTS5 全文索引(trigram tokenizer 支持中文搜索)
    CREATE VIRTUAL TABLE IF NOT EXISTS messages_fts USING fts5(
      content,
      content='messages',
      content_rowid='id',
      tokenize='trigram'
    )
  `)

  dbInstance.run(`
    -- FTS5 同步触发器
    CREATE TRIGGER IF NOT EXISTS messages_ai AFTER INSERT ON messages BEGIN
      INSERT INTO messages_fts(rowid, content) VALUES (new.id, new.content);
    END
  `)

  dbInstance.run(`
    CREATE TRIGGER IF NOT EXISTS messages_ad AFTER DELETE ON messages BEGIN
      INSERT INTO messages_fts(messages_fts, rowid, content) VALUES('delete', old.id, old.content);
    END
  `)

  dbInstance.run(`
    CREATE TRIGGER IF NOT EXISTS messages_au AFTER UPDATE ON messages BEGIN
      INSERT INTO messages_fts(messages_fts, rowid, content) VALUES('delete', old.id, old.content);
      INSERT INTO messages_fts(rowid, content) VALUES (new.id, new.content);
    END
  `)

  dbInstance.run(`
    -- 向量表(用于 sqlite-vec) - 注:Bun sqlite 暂不支持虚拟表,简化处理
    CREATE TABLE IF NOT EXISTS message_embeddings (
      id INTEGER PRIMARY KEY AUTOINCREMENT,
      msg_id INTEGER NOT NULL,
      embedding BLOB,
      text_hash TEXT,
      created_at TEXT DEFAULT (datetime('now')),
      UNIQUE(msg_id)
    )
  `)

  dbInstance.run(`
    -- 消息删除时清理向量
    CREATE TRIGGER IF NOT EXISTS messages_ad_delete_vec AFTER DELETE ON messages BEGIN
      DELETE FROM message_embeddings WHERE msg_id = old.id;
    END
  `)

  dbInstance.run(`
    -- 更新会话时间
    CREATE TRIGGER IF NOT EXISTS session_updated_after_insert AFTER INSERT ON messages BEGIN
      UPDATE sessions SET updated_at = datetime('now') WHERE id = new.session_id;
    END
  `)

  dbInstance.run(`
    CREATE TRIGGER IF NOT EXISTS session_updated_after_update AFTER UPDATE ON messages BEGIN
      UPDATE sessions SET updated_at = datetime('now') WHERE id = new.session_id;
    END
  `)

  dbInstance.run(`
    -- Embedding 失败记录表
    CREATE TABLE IF NOT EXISTS embedding_failures (
      msg_id INTEGER PRIMARY KEY,
      fail_count INTEGER DEFAULT 1,
      last_error TEXT,
      updated_at TEXT DEFAULT (datetime('now'))
    )
  `)

  dbInstance.run(`
    -- 索引
    CREATE INDEX IF NOT EXISTS idx_messages_session ON messages(session_id)
  `)

  dbInstance.run(`
    CREATE INDEX IF NOT EXISTS idx_messages_created ON messages(created_at)
  `)

  console.log('  ✅ ' + t('memory.db_done') + ': ' + DB_PATH)
}

/** 关闭数据库连接 - P3 #32: 标记关闭状态 */
export function closeDb(): void {
  if (db) {
    db.close()
    db = null
    isClosed = true
  }
}

export { DB_PATH }