dsh-diary-memory 0.1.1 → 0.1.2

package/src/session-memory.ts

@@ -0,0 +1,151 @@
+import { AsyncLocalStorage } from 'node:async_hooks'
+import { symbols } from '@deepseek-ai/cordis'
+import type { Context } from '@deepseek-ai/cordis'
+import type { Agent } from '@deepseek-ai/dsh-agent'
+import type { BasicCompactionEngine } from '@deepseek-ai/dsh-compaction-basic'
+import type {} from '@deepseek-ai/dsh-agent-preset-registry'
+import type { Message } from '@deepseek-ai/dsh-llm'
+import type { DiarySummary } from './diary.js'
+import { selectRoundRange } from './diary.js'
+import { appendRecalledContext, prepareDiarySummary } from './diary-runtime.js'
+import { DiaryOrchestrator } from './orchestrator.js'
+import type { RetrievalService } from './retrieval-service.js'
+import { registerRecall } from './tools.js'
+import { createScope, scopeOf } from '@deepseek-ai/dsh-scope'
+import type { Scope } from '@deepseek-ai/dsh-scope'
+
+type NativeEngine = Pick<BasicCompactionEngine, 'compactIfNeeded' | 'compactRegion' | 'compactNow'> & {
+  readonly ctx: Context
+  summarize(input: { readonly messages: readonly Message[] }, agent: Agent, signal?: AbortSignal): Promise<DiarySummary>
+}
+
+/** Route the existing native engine by Session, without rebinding its preset. */
+export class SessionMemory {
+  private readonly restorers = new Map<object, () => void>()
+  private readonly operations = new AsyncLocalStorage<{ agent: Agent; enabled: boolean }>()
+  private readonly jobs = new Set<Promise<unknown>>()
+  private readonly lifetime = new AbortController()
+  private closed = false
+  private readonly recallScopes = new Map<Agent, Scope>()
+
+  constructor(private readonly ctx: Context, private readonly memory: RetrievalService) {
+    ctx.effect(() => () => this.dispose(), 'diary-memory: existing sessions')
+    ctx.on('agent/created', async ({ agent }) => { this.attach(agent) })
+    ctx.on('agent/disposed', async ({ agent }) => {
+      const scope = this.recallScopes.get(agent)
+      this.recallScopes.delete(agent)
+      await scope?.dispose()
+    })
+    ctx.on('agent/pre-step', async ({ agent, messages, signal }, next) => {
+      const enabled = this.selected(agent)
+      const engine = this.attach(agent)
+      if (enabled && engine && !(engine instanceof DiaryOrchestrator) && !this.recallScopes.has(agent)) {
+        const key = scopeOf(agent.ctx)
+        if (!key) throw new Error('指定会话缺少原生 Agent scope')
+        const scope = createScope(ctx, key)
+        registerRecall(scope.ctx)
+        this.recallScopes.set(agent, scope)
+      }
+      if (!enabled) {
+        const scope = this.recallScopes.get(agent)
+        this.recallScopes.delete(agent)
+        await scope?.dispose()
+      }
+      const decision = await next()
+      // The companion preset already owns its recall middleware.
+      if (!enabled || engine instanceof DiaryOrchestrator) return decision
+      return appendRecalledContext(ctx, memory, decision, messages, signal)
+    }, { prepend: true })
+    for (const agent of ctx.agents.list()) this.attach(agent)
+  }
+
+  private selected(agent: Agent): boolean { return !this.closed && this.memory.config.sessionIds.includes(agent.session.id) }
+
+  attach(agent: Agent): BasicCompactionEngine | undefined {
+    const provider = this.ctx.agentPresets.serviceFor(agent, 'compaction')
+    if (!provider) {
+      if (this.selected(agent)) throw new Error('指定会话的预设没有原生压缩器，无法接入日记记忆')
+      return undefined
+    }
+    if (this.closed) return provider as BasicCompactionEngine
+    // Cordis proxies expose their original instance through this public symbol.
+    const engine = ((provider as unknown as Record<symbol, unknown>)[symbols.original] ?? provider) as NativeEngine
+    if (engine instanceof DiaryOrchestrator || this.restorers.has(engine)) return provider as BasicCompactionEngine
+    const descriptors = new Map<string, PropertyDescriptor | undefined>()
+    const replacements = new Map<string, unknown>()
+    const bridge = this
+    const enabled = (owner: Agent) => {
+      const operation = bridge.operations.getStore()
+      return operation?.agent === owner ? operation.enabled : bridge.selected(owner)
+    }
+    const replace = <K extends 'compactIfNeeded' | 'compactRegion' | 'compactNow' | 'summarize'>(name: K, method: NativeEngine[K]) => {
+      descriptors.set(name, Object.getOwnPropertyDescriptor(engine, name))
+      replacements.set(name, method)
+      Object.defineProperty(engine, name, { configurable: true, writable: true, value: method })
+    }
+    const summarize = engine.summarize
+    replace('summarize', function (this: NativeEngine, input, owner, signal) {
+      if (!enabled(owner)) return summarize.call(this, input, owner, signal)
+      return prepareDiarySummary(this.ctx, bridge.memory, input, owner, AbortSignal.any([bridge.lifetime.signal, ...signal ? [signal] : []]))
+    })
+    const compactIfNeeded = engine.compactIfNeeded
+    replace('compactIfNeeded', function (this: NativeEngine, owner, trigger, signal) {
+      const selected = enabled(owner)
+      return bridge.run(owner, selected, async () => {
+        if (!selected) return compactIfNeeded.call(this, owner, trigger, signal)
+        await bridge.memory.reconcile(owner.session)
+        signal.throwIfAborted()
+        const config = bridge.memory.config
+        if (trigger === 'pressure' && bridge.memory.stateOf(owner.session).completedRounds >= config.compressionRounds) {
+          const range = selectRoundRange(owner.session, config.retainRounds)
+          if (range) return this.compactRegion(range.start, range.end, owner, signal)
+        }
+        return compactIfNeeded.call(this, owner, trigger, signal)
+      })
+    })
+    const compactRegion = engine.compactRegion
+    replace('compactRegion', function (this: NativeEngine, start, end, owner, signal) {
+      const selected = enabled(owner)
+      return bridge.run(owner, selected, async () => {
+        if (!selected) return compactRegion.call(this, start, end, owner, signal)
+        try { return await compactRegion.call(this, start, end, owner, signal) }
+        finally { await bridge.memory.reconcile(owner.session) }
+      })
+    })
+    const compactNow = engine.compactNow
+    replace('compactNow', function (this: NativeEngine, owner, signal, sourceCommandId) {
+      const selected = enabled(owner)
+      return bridge.run(owner, selected, async () => {
+        if (!selected) return compactNow.call(this, owner, signal, sourceCommandId)
+        try { return await compactNow.call(this, owner, signal, sourceCommandId) }
+        finally { await bridge.memory.reconcile(owner.session) }
+      })
+    })
+    this.restorers.set(engine, () => {
+      for (const [name, descriptor] of descriptors) {
+        if (Object.getOwnPropertyDescriptor(engine, name)?.value !== replacements.get(name)) continue
+        if (descriptor) Object.defineProperty(engine, name, descriptor)
+        else Reflect.deleteProperty(engine, name)
+      }
+    })
+    return provider as BasicCompactionEngine
+  }
+
+  private run<T>(agent: Agent, enabled: boolean, operation: () => Promise<T>): Promise<T> {
+    const job = this.memory.track(() => this.operations.run({ agent, enabled }, operation))
+    this.jobs.add(job)
+    void job.then(() => this.jobs.delete(job), () => this.jobs.delete(job))
+    return job
+  }
+
+  async dispose(): Promise<void> {
+    if (this.closed) return
+    this.closed = true
+    this.lifetime.abort(new Error('指定会话的日记记忆已关闭'))
+    await Promise.allSettled([...this.jobs])
+    for (const restore of this.restorers.values()) restore()
+    this.restorers.clear()
+    await Promise.allSettled([...this.recallScopes.values()].map(scope => scope.dispose()))
+    this.recallScopes.clear()
+  }
+}

package/src/session-selection.ts

@@ -0,0 +1,4 @@
+export function validateSessionIds(value: unknown): string[] {
+  if (!Array.isArray(value) || value.length > 1000 || value.some(id => typeof id !== 'string' || !id.length || id.length > 256 || /[\r\n\0]/.test(id))) throw new Error('指定会话必须是至多 1000 个有效会话 ID 的数组')
+  return [...new Set(value as string[])]
+}

package/src/sqlite.ts

@@ -0,0 +1,78 @@
+import { createRequire } from 'node:module'
+import type { DatabaseSync, SQLInputValue } from 'node:sqlite'
+
+export type SqliteDriver = 'auto' | 'builtin' | 'native'
+export interface SqliteStatement {
+  run(...parameters: SQLInputValue[]): { changes: number | bigint; lastInsertRowid: number | bigint }
+  get(...parameters: SQLInputValue[]): Record<string, unknown> | undefined
+  all(...parameters: SQLInputValue[]): Record<string, unknown>[]
+  iterate(...parameters: SQLInputValue[]): IterableIterator<Record<string, unknown>>
+}
+export interface SqliteConnection {
+  readonly open: boolean
+  exec(sql: string): unknown
+  prepare(sql: string): SqliteStatement
+  pragma(sql: string, options?: { simple: boolean }): unknown
+  transaction<T>(operation: () => T): () => T
+  close(): unknown
+}
+
+const require = createRequire(import.meta.url)
+
+/** Use the same SQLite file format without requiring an Electron-specific addon. */
+function builtinDatabase(filename: string): SqliteConnection {
+  const { DatabaseSync: Database } = require('node:sqlite') as { DatabaseSync: typeof DatabaseSync }
+  const database = new Database(filename)
+  let open = true
+  let transactionId = 0
+  const normalize = (row: Record<string, unknown> | undefined) => row === undefined ? undefined
+    : Object.fromEntries(Object.entries(row).map(([key, value]) => [key, value instanceof Uint8Array ? Buffer.from(value) : value]))
+  return {
+    get open() { return open },
+    exec: sql => database.exec(sql),
+    prepare(sql) {
+      const statement = database.prepare(sql)
+      return {
+        run: (...parameters) => statement.run(...parameters),
+        get: (...parameters) => normalize(statement.get(...parameters)),
+        all: (...parameters) => statement.all(...parameters).map(row => normalize(row)!),
+        *iterate(...parameters) { for (const row of statement.iterate(...parameters)) yield normalize(row)! },
+      }
+    },
+    pragma(sql, options) {
+      const rows = database.prepare(`PRAGMA ${sql}`).all()
+      return options?.simple ? Object.values(rows[0] ?? {})[0] : rows
+    },
+    transaction: operation => () => {
+      // Savepoints also give nested synchronous operations rollback isolation.
+      const point = `diary_${++transactionId}`
+      database.exec(`SAVEPOINT ${point}`)
+      try {
+        const result = operation()
+        if (result instanceof Promise) throw new TypeError('SQLite 事务必须同步完成')
+        database.exec(`RELEASE ${point}`)
+        return result
+      } catch (error) {
+        database.exec(`ROLLBACK TO ${point}; RELEASE ${point}`)
+        throw error
+      }
+    },
+    close() { database.close(); open = false },
+  }
+}
+
+export function openDatabase(filename: string, driver: SqliteDriver = 'auto'): { db: SqliteConnection; backend: 'better-sqlite3' | 'node:sqlite' } {
+  if (driver !== 'builtin') {
+    try {
+      const Database = require('better-sqlite3') as typeof import('better-sqlite3')
+      return { db: new Database(filename) as unknown as SqliteConnection, backend: 'better-sqlite3' }
+    } catch (error) {
+      const failure = error as NodeJS.ErrnoException
+      const unavailable = failure.code === 'MODULE_NOT_FOUND' || failure.code === 'ERR_DLOPEN_FAILED'
+        || failure.message.startsWith('Could not locate the bindings file.')
+      // Disk, SQL, permission and corruption failures must never select another backend.
+      if (driver === 'native' || !unavailable) throw error
+    }
+  }
+  return { db: builtinDatabase(filename), backend: 'node:sqlite' }
+}

package/src/tools.ts

@@ -6,6 +6,10 @@ export const inject = ['tools', 'diaryMemory']
 
 /** Model-visible read-only recall, scoped by the companion preset and namespace. */
 export function apply(ctx: Context): void {
+  registerRecall(ctx)
+}
+
+export function registerRecall(ctx: Context): void {
   ctx.tools.register({
     name: 'diary_recall',
     description: '想起我们之前的经历、称呼、偏好、情绪变化、约定或我自己的反思。输入自然语言，会检索历史日记并按相关性重排。日记中的主观感受和不确定解释不等于已确认事实。',