dsh-account-pool 0.1.0

package/lib/usage.js

@@ -0,0 +1,420 @@
+/**
+ * 用量统计：逐请求记录 token 消耗，按 (时间片, 账号, 模型) 分桶累计。
+ *
+ * 统计口径参考 workbuddy2api 网关的 internal/usage：
+ *   - 按小时分桶；桶数超过上限时，把「超过保留期」的小时桶折叠为日桶
+ *     （注意：折叠由桶数上限触发，不是按时间定时触发——与参照实现一致）
+ *   - 记录请求数、失败数、prompt/completion/total tokens、延迟、吐字速率
+ *   - **失败也计入请求数**——否则「重试放大」在统计里看不见
+ *
+ * 与账号池里那个「每账号一个累计计数器」的区别：那边只有总量与最近一次，
+ * 没有时间维度，无法回答「今天各模型用了多少」这类问题。
+ *
+ * 落盘到 $DSH_HOME/.account-pool.usage.json，防抖写入，重启不丢。
+ */
+
+import { readFile, writeFile, rename, mkdir, stat } from 'node:fs/promises'
+import { dirname, join } from 'node:path'
+
+/**
+ * 小时桶的「保留期」判据：早于这个时长的小时桶在折叠时会被并成日桶。
+ *
+ * 注意它不是定时器——折叠只在桶数超过 MAX_BUCKETS 时触发，
+ * 这个常量决定「哪些桶算旧、可以并」。它也用来估算容量上界：
+ * 桶数 ≈ 账号数 × 模型数 × (保留期小时数 + 已过天数)。
+ */
+const HOURLY_KEEP_MS = 90 * 24 * 3600 * 1000
+
+/** 落盘防抖间隔。 */
+const FLUSH_INTERVAL_MS = 30_000
+
+/** 桶数硬上限：超过就立刻折叠一次，防止异常流量把内存撑爆。 */
+const MAX_BUCKETS = 200_000
+
+/** 落盘格式版本。 */
+const FILE_VERSION = 1
+
+/** 临时文件名序号：保证并发写入时各自的临时文件互不冲突。 */
+let tmpSeq = 0
+
+/** 本地时区的小时键。 */
+function hourKey(date) {
+  const pad = (n) => String(n).padStart(2, '0')
+  return `${date.getFullYear()}-${pad(date.getMonth() + 1)}-${pad(date.getDate())}T${pad(date.getHours())}`
+}
+
+/** 本地时区的日键。 */
+function dayKey(date) {
+  const pad = (n) => String(n).padStart(2, '0')
+  return `${date.getFullYear()}-${pad(date.getMonth() + 1)}-${pad(date.getDate())}`
+}
+
+/** 一个桶：某个 (时间片, 账号, 模型) 的累计量。 */
+function newBucket(scope, accountId, model) {
+  return {
+    scope, accountId, model,
+    requests: 0, errors: 0,
+    promptTokens: 0, completionTokens: 0, totalTokens: 0,
+    latencySumMs: 0, latencySamples: 0,
+    tpsSum: 0, tpsSamples: 0,
+    credits: 0,
+  }
+}
+
+/**
+ * 用量记录器。
+ *
+ * 用法：
+ *   const usage = new UsageRecorder({ file })
+ *   usage.add({ accountId, model, ok, usage: {...}, latencyMs })
+ *   usage.snapshot(72)   // 取聚合视图
+ */
+export class UsageRecorder {
+  /**
+   * @param {object} options
+   * @param {string} options.file 落盘路径
+   * @param {() => number} [options.now] 取当前时间（测试注入）
+   */
+  constructor({ file, now = () => Date.now(), legacyFiles = [] }) {
+    this.file = file
+    /**
+     * 历史文件名的只读迁移来源。
+     *
+     * 插件改名（workbuddy-pool → account-pool）后用量文件名也变了；
+     * 不做回落的话历史用量会"归零"。只读：一旦刷新就写到新文件。
+     */
+    this.legacyFiles = legacyFiles
+    this.now = now
+    /** key: `${scope}|${accountId}|${model}` → bucket */
+    this.buckets = new Map()
+    this.dirty = false
+    this.timer = null
+    this.loaded = false
+    /**
+     * 落盘文件的字节数。snapshot() 是同步的，没法现 stat（那是异步 API），
+     * 所以在 load/flush 时顺手量一次缓存起来。
+     */
+    this.fileBytes = 0
+    /** 是否有一次 flush 正在执行（防止两次写同时用同一个临时文件）。 */
+    this.flushing = false
+  }
+
+  /** 从磁盘加载历史数据；不存在或损坏时从空开始。 */
+  async load() {
+    if (this.loaded) return
+    this.loaded = true
+    try {
+      // 优先读当前文件名；读不到再依次试历史文件名（只读迁移）。
+      let source = this.file
+      // 顺手量一下文件大小，供 snapshot 展示（那里不能 await）。
+      this.fileBytes = await stat(source).then(info => info.size).catch(() => 0)
+      let raw
+      try {
+        raw = await readFile(source, 'utf8')
+      } catch {
+        for (const candidate of this.legacyFiles) {
+          try {
+            raw = await readFile(candidate, 'utf8')
+            source = candidate
+            this.fileBytes = await stat(candidate).then(info => info.size).catch(() => 0)
+            break
+          } catch {
+            // 该历史文件不存在，试下一个
+          }
+        }
+        if (raw === undefined) return
+      }
+      const doc = JSON.parse(raw)
+      if (doc?.version !== FILE_VERSION || !Array.isArray(doc.buckets)) return
+      for (const b of doc.buckets) {
+        if (typeof b?.scope !== 'string') continue
+        this.buckets.set(`${b.scope}|${b.accountId}|${b.model}`, {
+          ...newBucket(b.scope, b.accountId, b.model),
+          ...b,
+        })
+      }
+    } catch {
+      // 文件不存在或损坏：从空开始，不阻塞插件
+    }
+  }
+
+  /**
+   * 记一次请求尝试。
+   *
+   * @param {object} entry
+   * @param {string} entry.accountId 用了哪个账号
+   * @param {string} entry.model 请求的模型
+   * @param {boolean} entry.ok 这次尝试是否成功（拿到 usage 即视为成功）
+   * @param {object} [entry.usage] 上游返回的 usage 对象
+   * @param {number} [entry.latencyMs] 耗时
+   */
+  add({ accountId, model, ok, usage, latencyMs }) {
+    const at = new Date(this.now())
+    const scope = `h:${hourKey(at)}`
+    const account = accountId || '(unknown)'
+    const modelId = model || '(unknown)'
+    const key = `${scope}|${account}|${modelId}`
+
+    let bucket = this.buckets.get(key)
+    if (bucket === undefined) {
+      bucket = newBucket(scope, account, modelId)
+      this.buckets.set(key, bucket)
+    }
+
+    bucket.requests += 1
+    if (!ok) bucket.errors += 1
+
+    const prompt = Number(usage?.prompt_tokens)
+    const completion = Number(usage?.completion_tokens)
+    const total = Number(usage?.total_tokens)
+    const credit = Number(usage?.credit)
+
+    if (Number.isFinite(prompt)) bucket.promptTokens += prompt
+    if (Number.isFinite(completion)) bucket.completionTokens += completion
+    // 上游没给 total 时用 pt+ct 兜底，保证总量口径连续。
+    if (Number.isFinite(total)) bucket.totalTokens += total
+    else if (Number.isFinite(prompt) || Number.isFinite(completion)) {
+      bucket.totalTokens += (Number.isFinite(prompt) ? prompt : 0) + (Number.isFinite(completion) ? completion : 0)
+    }
+    if (Number.isFinite(credit)) bucket.credits += credit
+
+    if (Number.isFinite(latencyMs) && latencyMs > 0) {
+      bucket.latencySumMs += latencyMs
+      bucket.latencySamples += 1
+      // 吐字速率只有拿到 completion 才有意义
+      if (Number.isFinite(completion) && completion >= 0) {
+        bucket.tpsSum += (completion * 1000) / latencyMs
+        bucket.tpsSamples += 1
+      }
+    }
+
+    this.dirty = true
+    this.scheduleFlush()
+    if (this.buckets.size > MAX_BUCKETS) this.rollup()
+  }
+
+  /**
+   * 把「早于保留期」的小时桶折叠为日桶。幂等：先累加再删源桶。
+   *
+   * 调用时机：仅在桶数超过 MAX_BUCKETS 时（见 add）。所以短期内桶数远低于
+   * 上限时不会折叠——这与参照实现一致，靠上限保护内存而非定时整理。
+   */
+  rollup() {
+    const cutoff = this.now() - HOURLY_KEEP_MS
+    for (const [key, bucket] of [...this.buckets]) {
+      if (!bucket.scope.startsWith('h:')) continue
+      // scope 形如 "h:2026-09-18T15"，补上秒后按本地时间解析
+      const parsed = Date.parse(bucket.scope.slice(2) + ':00:00')
+      if (Number.isNaN(parsed) || parsed >= cutoff) continue
+
+      const dayScope = `d:${dayKey(new Date(parsed))}`
+      const dayKeyStr = `${dayScope}|${bucket.accountId}|${bucket.model}`
+      let target = this.buckets.get(dayKeyStr)
+      if (target === undefined) {
+        target = newBucket(dayScope, bucket.accountId, bucket.model)
+        this.buckets.set(dayKeyStr, target)
+      }
+      target.requests += bucket.requests
+      target.errors += bucket.errors
+      target.promptTokens += bucket.promptTokens
+      target.completionTokens += bucket.completionTokens
+      target.totalTokens += bucket.totalTokens
+      target.latencySumMs += bucket.latencySumMs
+      target.latencySamples += bucket.latencySamples
+      target.tpsSum += bucket.tpsSum
+      target.tpsSamples += bucket.tpsSamples
+      target.credits += bucket.credits
+      this.buckets.delete(key)
+      this.dirty = true
+    }
+  }
+
+  /** 安排一次防抖落盘。 */
+  scheduleFlush() {
+    if (this.timer !== null) return
+    this.timer = setTimeout(() => {
+      this.timer = null
+      void this.flush()
+    }, FLUSH_INTERVAL_MS)
+    // 定时器不该阻止进程退出
+    this.timer.unref?.()
+  }
+
+  /** 落盘（原子替换）。失败只记日志，不抛错。 */
+  async flush() {
+    // 已在写就跳过：dirty 标志本身不足以完全防重入
+    // （两次 add 穿插时仍可能同时进入），显式标志更稳。
+    if (this.flushing || !this.dirty) return
+    this.flushing = true
+    this.dirty = false
+    try {
+      await mkdir(dirname(this.file), { recursive: true })
+      const doc = {
+        version: FILE_VERSION,
+        saved: new Date(this.now()).toISOString(),
+        buckets: [...this.buckets.values()],
+      }
+      // 临时文件名带唯一后缀：即使有绕过 flushing 的路径也不会踩踏
+      const tmp = `${this.file}.${process.pid}.${++tmpSeq}.tmp`
+      const payload = JSON.stringify(doc)
+      await writeFile(tmp, payload, { encoding: 'utf8', mode: 0o600 })
+      await rename(tmp, this.file)
+      // 用实际写入的字节数（比再 stat 一次少一次系统调用）
+      this.fileBytes = Buffer.byteLength(payload, 'utf8')
+    } catch {
+      // 落盘失败不该影响请求处理；下次 add 会再触发
+      this.dirty = true
+    } finally {
+      this.flushing = false
+    }
+  }
+
+  /** 停止定时器并立即落盘（插件卸载时调用）。 */
+  async close() {
+    if (this.timer !== null) {
+      clearTimeout(this.timer)
+      this.timer = null
+    }
+    await this.flush()
+  }
+
+  /**
+   * 聚合出视图数据。
+   *
+   * @param {number} hours 时序返回多少个小时点（默认 72）
+   * @returns {{totals, byAccount, byModel, series, buckets, credits}}
+   */
+  snapshot(hours = 72) {
+    const window = hours > 0 && hours <= 24 * 60 ? hours : 72
+    const cutoff = this.now() - window * 3600 * 1000
+
+    const totals = newAcc()
+    const byAccount = new Map()
+    const byModel = new Map()
+    const series = []
+
+    for (const bucket of this.buckets.values()) {
+      const parsed = bucket.scope.startsWith('h:')
+        ? Date.parse(bucket.scope.slice(2) + ':00:00')
+        : Date.parse(bucket.scope.slice(2) + 'T00:00:00')
+
+      // 总量按全部桶算（含日桶），时序只取窗口内的点
+      addTo(totals, bucket)
+
+      const accountAcc = byAccount.get(bucket.accountId) ?? newAcc()
+      addTo(accountAcc, bucket)
+      byAccount.set(bucket.accountId, accountAcc)
+
+      const modelAcc = byModel.get(bucket.model) ?? newAcc()
+      addTo(modelAcc, bucket)
+      byModel.set(bucket.model, modelAcc)
+
+      if (!Number.isNaN(parsed) && parsed >= cutoff) {
+        series.push({
+          t: bucket.scope.slice(2),
+          kind: bucket.scope.startsWith('h:') ? 'hour' : 'day',
+          accountId: bucket.accountId,
+          model: bucket.model,
+          ...finish(bucketAccToAgg(bucket)),
+        })
+      }
+    }
+
+    series.sort((a, b) => a.t.localeCompare(b.t))
+
+    return {
+      totals: finish(totals),
+      byAccount: [...byAccount].map(([key, acc]) => ({ key, ...finish(acc) })).sort((a, b) => b.totalTokens - a.totalTokens),
+      byModel: [...byModel].map(([key, acc]) => ({ key, ...finish(acc) })).sort((a, b) => b.totalTokens - a.totalTokens),
+      series,
+      buckets: this.buckets.size,
+      fileBytes: this.fileBytes,
+      // 最早的分片即数据起点（只在有数据时给）
+      ...(series.length === 0 ? {} : { since: series[0].t }),
+    }
+  }
+}
+
+/** 新建一个累加器；均值需要样本数才能正确加权，不能对每桶均值再取平均。 */
+function newAcc() {
+  return {
+    requests: 0, errors: 0,
+    promptTokens: 0, completionTokens: 0, totalTokens: 0, credits: 0,
+    latencySumMs: 0, latencySamples: 0,
+    tpsSum: 0, tpsSamples: 0,
+  }
+}
+
+/** 一个桶直接转成累加器形状。 */
+function bucketAccToAgg(bucket) {
+  return {
+    requests: bucket.requests,
+    errors: bucket.errors,
+    promptTokens: bucket.promptTokens,
+    completionTokens: bucket.completionTokens,
+    totalTokens: bucket.totalTokens,
+    credits: bucket.credits,
+    latencySumMs: bucket.latencySumMs,
+    latencySamples: bucket.latencySamples,
+    tpsSum: bucket.tpsSum,
+    tpsSamples: bucket.tpsSamples,
+  }
+}
+
+function addTo(acc, bucket) {
+  acc.requests += bucket.requests
+  acc.errors += bucket.errors
+  acc.promptTokens += bucket.promptTokens
+  acc.completionTokens += bucket.completionTokens
+  acc.totalTokens += bucket.totalTokens
+  acc.credits += bucket.credits
+  acc.latencySumMs += bucket.latencySumMs
+  acc.latencySamples += bucket.latencySamples
+  acc.tpsSum += bucket.tpsSum
+  acc.tpsSamples += bucket.tpsSamples
+}
+
+/** 算出均值并去掉中间累加字段。 */
+function finish(acc) {
+  return {
+    requests: acc.requests,
+    errors: acc.errors,
+    promptTokens: acc.promptTokens,
+    completionTokens: acc.completionTokens,
+    totalTokens: acc.totalTokens,
+    credits: acc.credits,
+    avgLatencyMs: acc.latencySamples > 0 ? acc.latencySumMs / acc.latencySamples : 0,
+    avgTokensPerSecond: acc.tpsSamples > 0 ? acc.tpsSum / acc.tpsSamples : 0,
+    // 延迟与速率的样本数：它们只统计**成功**请求，可能远少于请求总数。
+    // 前端靠它区分「没有成功请求」（显示 —）与「就是很快」（显示真实值）。
+    latencySamples: acc.latencySamples,
+  }
+}
+
+/**
+ * 从 SSE 数据帧里提取 usage 与 model。
+ *
+ * 上游在流末尾发一帧带 usage 的数据（含 credit 积分消耗）。
+ * 这里只扫文本，不需要完整解析 SSE——目标字段明确且只出现在数据行。
+ *
+ * @param {string} text 一帧或一段 SSE 文本
+ * @returns {{usage?:object, model?:string}|undefined}
+ */
+export function extractUsage(text) {
+  if (typeof text !== 'string' || !text.includes('"usage"')) return undefined
+  for (const line of text.split('\n')) {
+    const trimmed = line.trim()
+    if (!trimmed.startsWith('data:')) continue
+    const payload = trimmed.slice(5).trim()
+    if (payload === '' || payload === '[DONE]') continue
+    try {
+      const frame = JSON.parse(payload)
+      if (frame?.usage && typeof frame.usage === 'object') {
+        return { usage: frame.usage, model: typeof frame.model === 'string' ? frame.model : undefined }
+      }
+    } catch {
+      // 不完整帧：忽略，后续 chunk 会补齐
+    }
+  }
+  return undefined
+}

package/package.json

@@ -0,0 +1,68 @@
+{
+  "name": "dsh-account-pool",
+  "displayName": "Account Pool",
+  "version": "0.1.0",
+  "description": "把多个 WorkBuddy / Trae 账号汇成账号池接入 DeepSeek Harness：自动切换、限流熔断退避、凭证安全落盘",
+  "type": "module",
+  "license": "MIT",
+  "repository": {
+    "type": "git",
+    "url": "git+https://github.com/cliii-one/dsh-account-pool.git"
+  },
+  "homepage": "https://github.com/cliii-one/dsh-account-pool#readme",
+  "bugs": {
+    "url": "https://github.com/cliii-one/dsh-account-pool/issues"
+  },
+  "author": "cliii-one",
+  "main": "lib/index.js",
+  "exports": {
+    ".": {
+      "default": "./lib/index.js"
+    },
+    "./client": "./lib/client.js",
+    "./cordis.patch.yml": "./cordis.patch.yml",
+    "./package.json": "./package.json"
+  },
+  "files": [
+    "lib",
+    "cordis.patch.yml",
+    "README.md"
+  ],
+  "engines": {
+    "node": "^22.19.0 || >=24.0.0"
+  },
+  "scripts": {
+    "check": "node scripts/check.mjs"
+  },
+  "dsh": {
+    "bundle": {
+      "patch": "./cordis.patch.yml"
+    },
+    "client": {
+      "platform": "web",
+      "inject": [
+        "@deepseek-ai/dsh-client-ui-slots"
+      ]
+    }
+  },
+  "keywords": [
+    "dsh",
+    "dsh-plugin",
+    "deepseek-harness",
+    "workbuddy",
+    "codebuddy",
+    "model-provider",
+    "account-pool",
+    "trae",
+    "trae-solo"
+  ],
+  "peerDependencies": {
+    "@deepseek-ai/cordis": ">=4.0.1 <5.0.0",
+    "@deepseek-ai/dsh-llm": "^0.1.6-alpha.2",
+    "@deepseek-ai/dsh-llm-pi-ai": "^0.1.6-alpha.2",
+    "@earendil-works/pi-ai": "^0.85.0"
+  },
+  "publishConfig": {
+    "access": "public"
+  }
+}