koishi-plugin-chat-patch 5.6.1 → 6.1.0

package/src/database.ts

@@ -1,17 +1,23 @@
-import { Level } from 'level'
-import { Context } from 'koishi'
-import path from 'node:path'
+import { Context } from 'koishi'
+import { mkdirSync } from 'node:fs'
+import path from 'node:path'
 import { createHash } from 'node:crypto'
 
 import { Config } from './config'
 import { ContactCacheItem, MessageRecord, PinnedState, SelfMessageRecord } from './types'
 import { PluginLogger } from './logger'
 
-function encodeKeyPart(value: string): string {
-  return encodeURIComponent(value)
-}
-
-function decodeKeyPart(value: string): string {
+function encodeKeyPart(value: string): string {
+  return encodeURIComponent(value)
+}
+
+function normalizeTimestampMs(value: unknown): number {
+  const num = Number(value ?? 0)
+  if (!Number.isFinite(num) || num <= 0) return 0
+  return num > 1e12 ? num : num * 1000
+}
+
+function decodeKeyPart(value: string): string {
   try {
     return decodeURIComponent(value)
   } catch {
@@ -27,8 +33,11 @@ function legacyMessagePrefix(platform: string, selfId: string, channelId: string
   return `m:${platform}:${selfId}:${channelId}:`
 }
 
-function messageKey(record: MessageRecord): string {
-  const time = String(record.timestamp).padStart(16, '0')
+function messageKey(record: MessageRecord): string {
+  // 统一用毫秒排序，前端 beforeTime 也传毫秒，避免按秒存储时分页取到同一批消息
+  const time = String(
+    normalizeTimestampMs(record.timestampMs ?? record.timestamp ?? record.receivedAt),
+  ).padStart(16, '0')
   return `${messagePrefix(record.platform, record.selfId, record.channelId || '')}${time}:${encodeKeyPart(record.id || 'unknown')}`
 }
 
@@ -37,8 +46,10 @@ function selfMessagePrefix(platform: string, selfId: string, channelId: string):
   return `sm:${encodeKeyPart(platform)}:${encodeKeyPart(selfId)}:${encodeKeyPart(channelId)}:`
 }
 
-function selfMessageKey(record: SelfMessageRecord): string {
-  const time = String(record.sentAt).padStart(16, '0')
+function selfMessageKey(record: SelfMessageRecord): string {
+  const time = String(
+    normalizeTimestampMs(record.timestampMs ?? record.timestamp ?? record.sentAt),
+  ).padStart(16, '0')
   return `${selfMessagePrefix(record.platform, record.selfId, record.channelId)}${time}:${encodeKeyPart(record.id || 'unknown')}`
 }
 
@@ -73,130 +84,240 @@ function isUsableGroupContact(item: ContactCacheItem): boolean {
   return true
 }
 
-type CompactableLevel = Level<string, string> & {
+type KvWriteOperation =
+  | { type: 'put'; key: string; value: string }
+  | { type: 'del'; key: string }
+
+interface KvIteratorOptions {
+  gte?: string
+  lte?: string
+  lt?: string
+  reverse?: boolean
+  limit?: number
+}
+
+interface LevelIteratorOptions {
+  gte?: string
+  lte?: string
+  lt?: string
+  reverse?: boolean
+  limit?: number
+}
+
+interface LevelDatabase {
+  open(): Promise<void>
+  close(): Promise<void>
+  put(key: string, value: string): Promise<void>
+  get(key: string): Promise<string>
+  del(key: string): Promise<void>
+  batch(operations: KvWriteOperation[]): Promise<void>
+  clear(): Promise<void>
+  iterator<K, V>(options?: LevelIteratorOptions): AsyncIterable<[K, V]>
   compactRange(start: string, end: string): Promise<void>
 }
 
-interface SharedDatabase {
-  dir: string
-  db: Level<string, string>
-  refs: number
-  opened: boolean
-  opening?: Promise<void>
-  closing?: Promise<void>
+interface LevelConstructor {
+  new (
+    location: string,
+    options?: { keyEncoding?: string; valueEncoding?: string },
+  ): LevelDatabase
+}
+
+interface LevelModule {
+  Level: LevelConstructor
+}
+
+interface WNodeService {
+  import<T>(packageName: string, options?: {
+    allowInstall?: boolean
+    useRequire?: boolean
+    version?: string
+  }): Promise<T>
 }
+
+interface KvStore {
+  readonly path: string
+  readonly driver: 'level'
+  readonly isOpen: boolean
+  put(key: string, value: string): Promise<void>
+  get(key: string): Promise<string>
+  del(key: string): Promise<void>
+  batch(operations: KvWriteOperation[]): Promise<void>
+  clear(): Promise<void>
+  vacuum(): Promise<void>
+  iterator(options?: KvIteratorOptions): AsyncGenerator<[string, string]>
+  close(): Promise<void> | void
+}
+
+// LevelDB 的读写性能适合这种消息缓存，恢复为之前的存储后端
+class LevelKV implements KvStore {
+  readonly driver = 'level' as const
+  private opened = false
 
-const sharedDatabases = new Map<string, SharedDatabase>()
+  constructor(
+    readonly path: string,
+    private readonly level: LevelDatabase,
+  ) {}
 
-// 同目录复用同一个 LevelDB 实例，避免 HMR 卸载/重载期间新旧插件抢占 LOCK
-async function acquireSharedDatabase(dir: string): Promise<SharedDatabase> {
-  let shared = sharedDatabases.get(dir)
-  if (shared?.closing) {
-    await shared.closing.catch(() => undefined)
-    shared = sharedDatabases.get(dir)
+  get isOpen(): boolean {
+    return this.opened
   }
-  if (!shared) {
-    shared = {
-      dir,
-      db: new Level<string, string>(dir, { keyEncoding: 'utf8', valueEncoding: 'utf8' }),
-      refs: 0,
-      opened: false,
-    }
-    sharedDatabases.set(dir, shared)
+
+  async open() {
+    await this.level.open()
+    this.opened = true
   }
-  shared.refs += 1
-  return shared
-}
 
-async function releaseSharedDatabase(shared: SharedDatabase): Promise<boolean> {
-  shared.refs -= 1
-  if (shared.refs > 0) return false
-  if (shared.closing) {
-    await shared.closing
-    return true
+  async close() {
+    if (!this.opened) return
+    await this.level.close()
+    this.opened = false
   }
-  if (!shared.opened) {
-    if (sharedDatabases.get(shared.dir) === shared) sharedDatabases.delete(shared.dir)
-    return true
+
+  async put(key: string, value: string) {
+    await this.level.put(key, value)
   }
-  const closing = shared.db.close().catch(() => undefined).finally(() => {
-    shared.opened = false
-    if (sharedDatabases.get(shared.dir) === shared) sharedDatabases.delete(shared.dir)
-  })
-  shared.closing = closing
-  await closing
-  return true
-}
 
-export class ChatDatabase {
-  private readonly dir: string
-  private shared?: SharedDatabase
+  async get(key: string): Promise<string> {
+    return await this.level.get(key)
+  }
 
-  constructor(
-    private ctx: Context,
-    private config: Config,
-    private logger: PluginLogger,
-  ) {
-    this.dir = path.resolve(ctx.baseDir, 'data', 'chat-patch', 'db')
+  async del(key: string) {
+    await this.level.del(key)
   }
 
-  private get db(): Level<string, string> {
-    if (!this.shared) throw new Error('LevelDB is not initialized')
-    return this.shared.db
+  async batch(operations: KvWriteOperation[]) {
+    await this.level.batch(operations)
   }
 
-  async initialize() {
-    this.shared = await acquireSharedDatabase(this.dir)
-    try {
-      await this.ensureOpen()
-      this.logger.logInfo('LevelDB 已打开:', this.shared.db.location)
-    } catch (error) {
-      await this.releaseShared()
-      throw error
-    }
+  async clear() {
+    await this.level.clear()
   }
 
-  async dispose() {
-    const closed = await this.releaseShared()
-    if (closed) this.logger.logInfo('LevelDB closed')
+  async vacuum() {
+    await this.level.compactRange('', '\uffff')
   }
 
-  private async ensureOpen() {
-    const shared = this.shared
-    if (!shared || shared.opened) return
-    if (shared.opening) {
-      await shared.opening
-      return
+  async *iterator(options: KvIteratorOptions = {}): AsyncGenerator<[string, string]> {
+    for await (const [key, value] of this.level.iterator<string, string>(options)) {
+      yield [key, value]
     }
-    const opening = shared.db.open().then(() => {
-      shared.opened = true
+  }
+}
+
+interface SharedDatabase {
+  path: string
+  db: KvStore
+  refs: number
+  closing?: Promise<void>
+}
+
+const sharedDatabases = new Map<string, SharedDatabase>()
+
+// 同文件复用同一个 LevelDB 连接，避免 HMR 卸载/重载期间互相持有旧句柄
+async function acquireSharedDatabase(filePath: string, Level: LevelConstructor): Promise<SharedDatabase> {
+  let shared = sharedDatabases.get(filePath)
+  if (shared?.closing) {
+    await shared.closing.catch(() => undefined)
+    shared = sharedDatabases.get(filePath)
+  }
+  if (!shared) {
+    mkdirSync(path.dirname(filePath), { recursive: true })
+    const level = new Level(filePath, {
+      keyEncoding: 'utf8',
+      valueEncoding: 'utf8',
     })
-    shared.opening = opening
-    try {
-      await opening
-    } finally {
-      shared.opening = undefined
-    }
+    const db = new LevelKV(filePath, level)
+    await db.open()
+    shared = {
+      path: filePath,
+      db,
+      refs: 0,
+    }
+    sharedDatabases.set(filePath, shared)
+  }
+  shared.refs += 1
+  return shared
+}
+
+async function releaseSharedDatabase(shared: SharedDatabase): Promise<boolean> {
+  shared.refs -= 1
+  if (shared.refs > 0) return false
+  if (shared.closing) {
+    await shared.closing
+    return true
+  }
+  const closing = Promise.resolve().then(() => {
+    if (shared.db.isOpen) return shared.db.close()
+  }).catch(() => undefined).finally(() => {
+    if (sharedDatabases.get(shared.path) === shared) sharedDatabases.delete(shared.path)
+  })
+  shared.closing = closing
+  await closing
+  return true
+}
+
+export class ChatDatabase {
+  private readonly dir: string
+  private shared?: SharedDatabase
+
+  constructor(
+    private ctx: Context,
+    private config: Config,
+    private logger: PluginLogger,
+  ) {
+    this.dir = path.resolve(ctx.baseDir, 'data', 'chat-patch', 'db')
+  }
+
+  private get db(): KvStore {
+    if (!this.shared) throw new Error('Store is not initialized')
+    return this.shared.db
+  }
+
+  async initialize() {
+    const levelModule = await this.loadLevelModule()
+    this.shared = await acquireSharedDatabase(this.dir, levelModule.Level)
+    this.logger.logInfo('LevelDB 已打开:', this.shared.db.driver, this.shared.path)
   }
 
-  private async releaseShared(): Promise<boolean> {
-    const shared = this.shared
-    this.shared = undefined
-    if (shared) return releaseSharedDatabase(shared)
-    return false
+  // 通过 w-node 服务动态安装并加载 level，原生 .node 不会进入插件依赖目录
+  private async loadLevelModule(): Promise<LevelModule> {
+    const nodeService = this.ctx.node
+    if (!nodeService) {
+      throw new Error('未检测到 w-node 服务，请先安装 koishi-plugin-w-node')
+    }
+    const levelModule = await nodeService.import<LevelModule>('level', {
+      version: '^10.0.0',
+      useRequire: true,
+    })
+    if (!levelModule?.Level) {
+      throw new Error('w-node 动态加载 level 失败')
+    }
+    return levelModule
   }
 
+  async dispose() {
+    const closed = await this.releaseShared()
+    if (closed) this.logger.logInfo('存储已关闭')
+  }
+
+  private async releaseShared(): Promise<boolean> {
+    const shared = this.shared
+    this.shared = undefined
+    if (shared) return releaseSharedDatabase(shared)
+    return false
+  }
+
   async clearAll() {
-    // 先清空记录，再强制压缩，让旧 .ldb 文件也能被回收
+    // 先清空记录，再强制压缩，让旧页也能被回收
     await this.db.clear()
-    const db = this.db as unknown as CompactableLevel
-    await db.compactRange('', '\uffff')
+    await this.db.vacuum()
     this.logger.logInfo('数据库缓存已全部清空并完成压缩')
   }
 
   async appendMessage(record: MessageRecord) {
     await this.db.put(messageKey(record), JSON.stringify(record))
-    await this.trimMessages(record.platform, record.selfId, record.channelId || '')
+    await this.trimChannel(record.platform, record.selfId, record.channelId || '')
   }
 
   async upsertSelfMessage(record: SelfMessageRecord) {
@@ -210,7 +331,7 @@ export class ChatDatabase {
       )
     }
     await this.db.put(selfMessageKey(record), JSON.stringify(record))
-    await this.trimSelfMessages(record.platform, record.selfId, record.channelId)
+    await this.trimChannel(record.platform, record.selfId, record.channelId)
   }
 
   async listSelfMessages(
@@ -221,7 +342,7 @@ export class ChatDatabase {
   ): Promise<SelfMessageRecord[]> {
     const result: SelfMessageRecord[] = []
     const prefix = selfMessagePrefix(platform, selfId, channelId)
-    for await (const [, value] of this.db.iterator<string, string>({
+    for await (const [, value] of this.db.iterator({
       gte: prefix,
       lte: `${prefix}\uffff`,
       reverse: true,
@@ -243,7 +364,7 @@ export class ChatDatabase {
     const result: SelfMessageRecord[] = []
     const prefix = selfMessagePrefix(platform, selfId, channelId)
     const before = `${prefix}${String(beforeTime).padStart(16, '0')}`
-    for await (const [, value] of this.db.iterator<string, string>({
+    for await (const [, value] of this.db.iterator({
       gte: prefix,
       lt: before,
       reverse: true,
@@ -263,7 +384,7 @@ export class ChatDatabase {
     patch: Partial<SelfMessageRecord>,
   ): Promise<boolean> {
     const prefix = `sm:${encodeKeyPart(platform)}:${encodeKeyPart(selfId)}:`
-    for await (const [key, value] of this.db.iterator<string, string>({
+    for await (const [key, value] of this.db.iterator({
       gte: prefix,
       lte: `${prefix}\uffff`,
     })) {
@@ -290,7 +411,7 @@ export class ChatDatabase {
   ): Promise<boolean> {
     const operations: Array<{ type: 'put'; key: string; value: string }> = []
     for (const prefix of [messagePrefix(platform, selfId, channelId), legacyMessagePrefix(platform, selfId, channelId)]) {
-      for await (const [key, value] of this.db.iterator<string, string>({
+      for await (const [key, value] of this.db.iterator({
         gte: prefix,
         lte: `${prefix}\uffff`,
       })) {
@@ -330,7 +451,7 @@ export class ChatDatabase {
       ? [selfMessagePrefix(platform, selfId, channelId)]
       : [`sm:${encodeKeyPart(platform)}:${encodeKeyPart(selfId)}:`]
     for (const prefix of prefixes) {
-      for await (const [, value] of this.db.iterator<string, string>({
+      for await (const [, value] of this.db.iterator({
         gte: prefix,
         lte: `${prefix}\uffff`,
       })) {
@@ -355,7 +476,7 @@ export class ChatDatabase {
   ) {
     const prefix = `sm:${encodeKeyPart(platform)}:${encodeKeyPart(selfId)}:`
     const toDelete: string[] = []
-    for await (const [key, value] of this.db.iterator<string, string>({
+    for await (const [key, value] of this.db.iterator({
       gte: prefix,
       lte: `${prefix}\uffff`,
     })) {
@@ -396,7 +517,7 @@ export class ChatDatabase {
   ): Promise<MessageRecord[]> {
     const result: MessageRecord[] = []
     for (const prefix of [messagePrefix(platform, selfId, channelId), legacyMessagePrefix(platform, selfId, channelId)]) {
-      for await (const [, value] of this.db.iterator<string, string>({
+      for await (const [, value] of this.db.iterator({
         gte: prefix,
         lte: `${prefix}\uffff`,
         reverse: true,
@@ -426,7 +547,7 @@ export class ChatDatabase {
     const result: MessageRecord[] = []
     for (const prefix of [messagePrefix(platform, selfId, channelId), legacyMessagePrefix(platform, selfId, channelId)]) {
       const before = `${prefix}${String(beforeTime).padStart(16, '0')}`
-      for await (const [, value] of this.db.iterator<string, string>({
+      for await (const [, value] of this.db.iterator({
         gte: prefix,
         lt: before,
         reverse: true,
@@ -449,7 +570,7 @@ export class ChatDatabase {
   async clearChannel(platform: string, selfId: string, channelId: string) {
     const operations: Array<{ type: 'del'; key: string }> = []
     for (const prefix of [messagePrefix(platform, selfId, channelId), legacyMessagePrefix(platform, selfId, channelId)]) {
-      for await (const [key] of this.db.iterator<string, string>({
+      for await (const [key] of this.db.iterator({
         gte: prefix,
         lte: `${prefix}\uffff`,
       })) {
@@ -457,7 +578,7 @@ export class ChatDatabase {
       }
     }
     const selfPrefix = selfMessagePrefix(platform, selfId, channelId)
-    for await (const [key] of this.db.iterator<string, string>({
+    for await (const [key] of this.db.iterator({
       gte: selfPrefix,
       lte: `${selfPrefix}\uffff`,
     })) {
@@ -521,7 +642,7 @@ export class ChatDatabase {
 
   async getAllMedia(): Promise<Array<{ filePath: string; channelId: string }>> {
     const result: Array<{ filePath: string; channelId: string }> = []
-    for await (const [, value] of this.db.iterator<string, string>({
+    for await (const [, value] of this.db.iterator({
       gte: 'media:',
       lte: 'media:\uffff',
     })) {
@@ -548,7 +669,7 @@ export class ChatDatabase {
   async removeMediaByPath(filePath: string) {
     const normalized = path.normalize(filePath)
     const toDelete: string[] = []
-    for await (const [key, value] of this.db.iterator<string, string>({
+    for await (const [key, value] of this.db.iterator({
       gte: 'media:',
       lte: 'media:\uffff',
     })) {
@@ -690,7 +811,7 @@ export class ChatDatabase {
       kind: 'new' | 'legacy'
       entry: ContactEntry
     }>()
-    for await (const [key, value] of this.db.iterator<string, string>({
+    for await (const [key, value] of this.db.iterator({
       gte: 'c:',
       lte: 'c:\uffff',
     })) {
@@ -742,36 +863,89 @@ export class ChatDatabase {
     return result
   }
 
-  private async trimMessages(platform: string, selfId: string, channelId: string) {
-    let count = 0
-    const toDelete: string[] = []
-    for (const prefix of [messagePrefix(platform, selfId, channelId), legacyMessagePrefix(platform, selfId, channelId)]) {
-      for await (const [key] of this.db.iterator<string, string>({
+  // 收到消息或机器人发送消息后统一裁剪，m 和 sm 合起来按时间保留上限
+  private async trimChannel(platform: string, selfId: string, channelId: string) {
+    const rows: Array<{ key: string; time: number }> = []
+    const prefixes = [
+      messagePrefix(platform, selfId, channelId),
+      legacyMessagePrefix(platform, selfId, channelId),
+      selfMessagePrefix(platform, selfId, channelId),
+    ]
+    for (const prefix of prefixes) {
+      for await (const [key, value] of this.db.iterator({
         gte: prefix,
         lte: `${prefix}\uffff`,
       })) {
-        count += 1
-        if (count > this.config.maxMessagesPerChannel) toDelete.push(key)
+        rows.push({ key, time: this.extractRecordTime(value) })
       }
     }
-    if (!toDelete.length) return
-    await this.db.batch(toDelete.map((key) => ({ type: 'del', key })))
-    this.logger.logInfo(`频道历史已裁剪 ${toDelete.length} 条`)
+    await this.trimRows(rows)
   }
 
-  private async trimSelfMessages(platform: string, selfId: string, channelId: string) {
-    let count = 0
-    const toDelete: string[] = []
-    const prefix = selfMessagePrefix(platform, selfId, channelId)
-    for await (const [key] of this.db.iterator<string, string>({
-      gte: prefix,
-      lte: `${prefix}\uffff`,
-    })) {
-      count += 1
-      if (count > this.config.maxMessagesPerChannel) toDelete.push(key)
+  // 启动时清理历史遗留的超量消息，避免配置项只在写入时才生效
+  async cleanupExcess() {
+    const groups = new Map<string, {
+      platform: string
+      selfId: string
+      channelId: string
+      rows: Array<{ key: string; time: number }>
+    }>()
+    const collect = async (prefix: string) => {
+      for await (const [key, value] of this.db.iterator({
+        gte: prefix,
+        lte: `${prefix}\uffff`,
+      })) {
+        let parsed: Record<string, unknown>
+        try {
+          parsed = JSON.parse(value) as Record<string, unknown>
+        } catch {
+          continue
+        }
+        const platform = typeof parsed.platform === 'string' ? parsed.platform : ''
+        const selfId = typeof parsed.selfId === 'string' ? parsed.selfId : ''
+        const channelId = typeof parsed.channelId === 'string' ? parsed.channelId : ''
+        if (!platform || !selfId) continue
+        const groupKey = JSON.stringify([platform, selfId, channelId])
+        let group = groups.get(groupKey)
+        if (!group) {
+          group = { platform, selfId, channelId, rows: [] }
+          groups.set(groupKey, group)
+        }
+        group.rows.push({ key, time: this.extractRecordTime(value) })
+      }
     }
-    if (!toDelete.length) return
-    await this.db.batch(toDelete.map((key) => ({ type: 'del', key })))
-    this.logger.logInfo(`机器人消息已裁剪 ${toDelete.length} 条`)
+    await collect('m:')
+    await collect('sm:')
+
+    let removed = 0
+    for (const group of groups.values()) {
+      const excess = group.rows.length - this.config.maxMessagesPerChannel
+      if (excess <= 0) continue
+      group.rows.sort((a, b) => a.time - b.time || (a.key < b.key ? -1 : 1))
+      const keys = group.rows.slice(0, excess).map((item) => item.key)
+      await this.db.batch(keys.map((key) => ({ type: 'del', key })))
+      removed += keys.length
+    }
+    if (removed) {
+      this.logger.logInfo(`启动时已裁剪 ${removed} 条超出上限的历史消息`)
+    }
+  }
+
+  private async trimRows(rows: Array<{ key: string; time: number }>) {
+    const excess = rows.length - this.config.maxMessagesPerChannel
+    if (excess <= 0) return
+    rows.sort((a, b) => a.time - b.time || (a.key < b.key ? -1 : 1))
+    const keys = rows.slice(0, excess).map((item) => item.key)
+    await this.db.batch(keys.map((key) => ({ type: 'del', key })))
+    this.logger.logInfo(`频道历史已裁剪 ${keys.length} 条`)
+  }
+
+  private extractRecordTime(value: string): number {
+    try {
+      const parsed = JSON.parse(value) as Record<string, unknown>
+      return normalizeTimestampMs(parsed.timestampMs ?? parsed.timestamp ?? parsed.sentAt ?? parsed.receivedAt)
+    } catch {
+      return 0
+    }
   }
 }