@chenmiao8563/dsh-token-ledger 0.1.0

package/lib/index.js

@@ -0,0 +1,253 @@
+/**
+ * `dsh-token-ledger` host half.
+ *
+ * A transparent, restart-safe token ledger for DeepSeek Harness. It folds the
+ * durable session event stream into per-session, per-day and per-model token
+ * totals, keeps the result in `<DSH_HOME>/token-ledger/ledger.json`, and
+ * answers `/tokens` in the conversation input.
+ *
+ * Design constraints, in order of importance:
+ *
+ * 1. **Installable anywhere.** The module imports nothing but Node builtins and
+ *    its own files, so it cannot fail on a missing or drifting peer dependency.
+ *    It also declares no install scripts, so `dsh plugin add` from a git URL
+ *    works without pre-authorizing a pnpm build.
+ * 2. **Never in the model's way.** It registers no prompt section, no message
+ *    and no model-facing tool, so it cannot change a request prefix and cannot
+ *    affect KV-cache reuse. The ledger is read by a command and by the CLI.
+ * 3. **Never the reason a session fails.** Every hook is wrapped so a ledger
+ *    fault degrades to a warning; folding is pure and cannot mutate session
+ *    state.
+ *
+ * @module dsh-token-ledger
+ */
+
+import { UsageLedger, inheritedCut } from './ledger.js'
+import { ledgerPaths, loadLedger, saveLedger, writeFileAtomic } from './store.js'
+import { join } from 'node:path'
+
+/** Plugin name, as it appears in the Cordis tree. */
+export const name = 'token-ledger'
+
+/** Milliseconds to coalesce ledger writes after activity. */
+const PERSIST_DEBOUNCE_MS = 2000
+
+/**
+ * Mount the ledger.
+ *
+ * @param {object} ctx - the Cordis context for this plugin's fiber.
+ * @param {{ ledgerPath?: string, backfill?: boolean }} [config] - composition config.
+ * @returns {void}
+ */
+export function apply(ctx, config = {}) {
+  const paths = ledgerPaths()
+  const ledgerPath = typeof config.ledgerPath === 'string' && config.ledgerPath !== '' ? config.ledgerPath : paths.ledger
+  const shouldBackfill = config.backfill !== false
+
+  const ledger = new UsageLedger()
+  const restored = ledger.restore(loadLedger(ledgerPath))
+
+  const log = {
+    /** @param {string} message @param {...unknown} rest */
+    info: (message, ...rest) => ctx.logger?.info?.(message, ...rest),
+    warn: (message, ...rest) => ctx.logger?.warn?.(message, ...rest),
+  }
+
+  let timer
+  let writing = Promise.resolve()
+
+  /**
+   * Write the ledger, serializing writes so two flushes cannot interleave.
+   *
+   * @returns {Promise<void>} settles when the file is on disk.
+   */
+  const flush = () => {
+    writing = writing
+      .then(() => {
+        saveLedger(ledgerPath, ledger.snapshot())
+      })
+      .catch((error) => {
+        log.warn('[token-ledger] could not write %s: %o', ledgerPath, error)
+      })
+    return writing
+  }
+
+  /** @param {number} [delay] - coalescing delay in milliseconds. */
+  const schedule = (delay = PERSIST_DEBOUNCE_MS) => {
+    if (timer !== undefined) clearTimeout(timer)
+    timer = setTimeout(() => {
+      timer = undefined
+      void flush()
+    }, delay)
+    // Do not hold the host process open just to write a ledger.
+    timer.unref?.()
+  }
+
+  /**
+   * Fold a live session's unconsumed events.
+   *
+   * @param {object} session - a live Session.
+   * @returns {void}
+   */
+  const adopt = (session) => {
+    try {
+      ledger.adoptSession(session)
+    } catch (error) {
+      log.warn('[token-ledger] could not fold session: %o', error)
+    }
+  }
+
+  // Sessions already live when this plugin mounts (a profile reload, or a
+  // ledger installed into a running host).
+  try {
+    const sessions = ctx.get?.('sessions')
+    if (sessions !== undefined && typeof sessions.list === 'function') {
+      for (const session of sessions.list()) adopt(session)
+    }
+  } catch (error) {
+    log.warn('[token-ledger] could not enumerate live sessions: %o', error)
+  }
+
+  /**
+   * Fold every session storage knows about, including ones from previous runs.
+   *
+   * This is what makes a fresh install immediately show real history instead of
+   * starting from zero. It is safe to repeat: a session already consumed is
+   * skipped by its cursor.
+   *
+   * @param {object} persistence - the `sessionPersistence` service.
+   * @returns {Promise<{ sessions: number, errors: number }>} the backfill tally.
+   */
+  const backfill = async (persistence) => {
+    let sessions = 0
+    let errors = 0
+    const headers = await persistence.list()
+    for (const header of headers) {
+      const id = String(header?.id ?? '')
+      if (id === '') continue
+      try {
+        const read = typeof persistence.inspect === 'function' ? persistence.inspect : persistence.load
+        const stored = await read.call(persistence, id)
+        const events = stored?.events
+        if (!Array.isArray(events)) continue
+        const meta = stored?.meta ?? header
+        ledger.adoptHistory({
+          sessionId: id,
+          events,
+          // Only a fork's prefix belongs to another session; see inheritedCut.
+          inheritedEventCount: inheritedCut(meta, stored?.inheritedEventCount ?? meta?.seedLength),
+        })
+        sessions += 1
+      } catch (error) {
+        errors += 1
+        log.warn('[token-ledger] could not read stored session %s: %o', id, error)
+      }
+    }
+    return { sessions, errors }
+  }
+
+  if (shouldBackfill) {
+    const onPersistence = (persistenceCtx) => {
+      void backfill(persistenceCtx.sessionPersistence)
+        .then((tally) => {
+          log.info(
+            '[token-ledger] backfilled %d stored session(s), %d unreadable; total %d tokens over %d calls',
+            tally.sessions,
+            tally.errors,
+            ledger.totals.totalTokens,
+            ledger.calls,
+          )
+          schedule(0)
+        })
+        .catch((error) => {
+          log.warn('[token-ledger] backfill failed: %o', error)
+        })
+    }
+    if (typeof ctx.inject === 'function') ctx.inject(['sessionPersistence'], onPersistence)
+  }
+
+  if (typeof ctx.on === 'function') {
+    ctx.on('session/created', (session) => {
+      adopt(session)
+      schedule()
+    })
+    ctx.on('session/event', (session) => {
+      adopt(session)
+      schedule()
+    })
+    ctx.on('session/disposed', (session) => {
+      adopt(session)
+      schedule(0)
+    })
+    ctx.on('session/end-seed', () => schedule())
+  }
+
+  if (!restored) log.info('[token-ledger] started a new ledger at %s', ledgerPath)
+
+  /**
+   * Handle `/tokens [summary|export|json|path]`.
+   *
+   * @param {{ rawInput?: string }} input - the command invocation.
+   * @returns {{ kind: 'success'|'error', text: string }} the command result.
+   */
+  const runCommand = ({ rawInput = '' } = {}) => {
+    const argument = String(rawInput).trim().split(/\s+/)[0]?.toLowerCase() ?? ''
+    switch (argument) {
+      case '':
+      case 'summary': {
+        return { kind: 'success', text: ledger.format() }
+      }
+      case 'path': {
+        return { kind: 'success', text: ledgerPath }
+      }
+      case 'json': {
+        return { kind: 'success', text: JSON.stringify(ledger.snapshot(), null, 2) }
+      }
+      case 'export': {
+        const stamp = new Date().toISOString().slice(0, 10)
+        const written = []
+        try {
+          written.push(writeFileAtomic(join(paths.exportsDir, `daily-${stamp}.csv`), ledger.toCsv('daily')))
+          written.push(writeFileAtomic(join(paths.exportsDir, `sessions-${stamp}.csv`), ledger.toCsv('sessions')))
+          written.push(writeFileAtomic(join(paths.exportsDir, `models-${stamp}.csv`), ledger.toCsv('models')))
+          written.push(saveLedger(join(paths.exportsDir, `ledger-${stamp}.json`), ledger.snapshot()))
+        } catch (error) {
+          return { kind: 'error', text: `export failed: ${error instanceof Error ? error.message : String(error)}` }
+        }
+        return { kind: 'success', text: `exported:\n${written.map((path) => `  ${path}`).join('\n')}` }
+      }
+      default: {
+        return {
+          kind: 'error',
+          text: 'usage: /tokens [summary|export|json|path]',
+        }
+      }
+    }
+  }
+
+  if (typeof ctx.inject === 'function') {
+    ctx.inject(['commands'], (commandCtx) => {
+      commandCtx.commands.register({
+        name: 'tokens',
+        description: 'Show the cumulative token ledger (per day and per model)',
+        input: 'summary | export | json | path',
+        handler: runCommand,
+      })
+    })
+  }
+
+  if (typeof ctx.effect === 'function') {
+    ctx.effect(
+      () => () => {
+        if (timer !== undefined) {
+          clearTimeout(timer)
+          timer = undefined
+        }
+        return flush()
+      },
+      'token-ledger: flush on dispose',
+    )
+  }
+}
+
+export default { name, apply }