dsh-modelscope-pool 0.0.0-stage → 0.8.0

package/lib/pool.js

@@ -0,0 +1,336 @@
+// 多 Key 池：轮换、失败冷却、到点自动恢复。
+// 只做内存状态与本地文件读写，不联网。
+import { mkdir, readFile, writeFile, chmod, stat } from 'node:fs/promises'
+import { join } from 'node:path'
+
+const COOLDOWN_MS = 5 * 60_000
+const CN_OFFSET_MS = 8 * 3600_000
+const DAY_MS = 86_400_000
+
+/**
+ * 额度用尽后的统一恢复时刻：北京时间最近的一个 00:01:00。
+ * 魔搭的额度按自然日（北京时间）重置，留一分钟余量，
+ * 避免刚好卡在服务端切换的瞬间又失败一轮。
+ */
+export function nextQuotaResetAt(now = Date.now()) {
+  const cn = now + CN_OFFSET_MS
+  const dayStart = Math.floor(cn / DAY_MS) * DAY_MS
+  let reset = dayStart + 60_000
+  if (reset <= cn) reset = dayStart + DAY_MS + 60_000
+  return reset - CN_OFFSET_MS
+}
+
+/** 掩码：永远不把完整 Key 交给界面、日志或响应头。 */
+export function maskKey(key) {
+  const s = String(key ?? '')
+  if (s.length <= 6) return '••••'
+  return `${s.slice(0, 3)}…${s.slice(-4)}`
+}
+
+export class KeyPool {
+  constructor({ dir, cooldownMs = COOLDOWN_MS, logger }) {
+    this.dir = dir
+    this.keysFile = join(dir, 'keys.json')
+    this.coolFile = join(dir, 'cooldowns.json')
+    this.cooldownMs = cooldownMs
+    this.logger = logger
+    this.keys = []
+    this.states = new Map() // id -> { until, fails, lastReason, lastError, used, ok }
+    this.cursor = 0
+    this.stats = { requests: 0, ok: 0, failed: 0, failovers: 0 }
+    this.onChange = null
+    this._lastCooldownCount = 0
+  }
+
+  async init() {
+    await mkdir(this.dir, { recursive: true, mode: 0o700 })
+    await this.reload()
+    await this._loadCool()
+    return this
+  }
+
+  /**
+   * 冷却状态落盘：额度用尽的 Key 要停到次日 00:01，
+   * 引擎重启（App 每次冷启动都会重启）后不能又拿它去撞墙。
+   */
+  async _loadCool() {
+    try {
+      const raw = JSON.parse(await readFile(this.coolFile, 'utf8'))
+      const now = Date.now()
+      for (const [id, v] of Object.entries(raw ?? {})) {
+        const st = this.states.get(id)
+        const until = Number(v?.until) || 0
+        if (!st || until <= now) continue
+        st.until = until
+        st.lastReason = String(v?.reason ?? '额度用尽').slice(0, 24)
+      }
+    } catch {}
+  }
+
+  async _saveCool() {
+    const now = Date.now()
+    const out = {}
+    for (const [id, st] of this.states) {
+      if (st.until > now) out[id] = { until: st.until, reason: st.lastReason }
+    }
+    try {
+      await writeFile(this.coolFile, JSON.stringify(out, null, 1) + '\n', { mode: 0o600 })
+    } catch {}
+  }
+
+  async reload() {
+    try {
+      const raw = JSON.parse(await readFile(this.keysFile, 'utf8'))
+      const list = Array.isArray(raw?.keys) ? raw.keys : []
+      this.keys = list
+        .filter((k) => k && typeof k.key === 'string' && k.key.length >= 8)
+        .map((k, i) => ({
+          id: typeof k.id === 'string' && k.id ? k.id : `k${i + 1}`,
+          key: k.key,
+          label: typeof k.label === 'string' ? k.label.slice(0, 24) : '',
+          region: k.region === 'ai' || k.region === 'cn' ? k.region : '',
+          addedAt: Number(k.addedAt) || 0,
+        }))
+    } catch {
+      this.keys = []
+    }
+    this.keys.forEach((k) => {
+      if (!this.states.has(k.id)) this.states.set(k.id, { until: 0, fails: 0, lastReason: '', lastError: '', used: 0, ok: 0 })
+    })
+    for (const id of [...this.states.keys()]) if (!this.keys.some((k) => k.id === id)) this.states.delete(id)
+    return this.keys.length
+  }
+
+  async _save() {
+    const body = JSON.stringify({ keys: this.keys }, null, 2) + '\n'
+    await writeFile(this.keysFile, body, { mode: 0o600 })
+    try {
+      await chmod(this.keysFile, 0o600)
+    } catch {}
+  }
+
+  has(key) {
+    return this.keys.some((k) => k.key === key)
+  }
+
+  /** 认出 Key 的版本后写回 keys.json；版本变了（比如换站点后）也能改。 */
+  async setRegion(id, region) {
+    const key = this.keys.find((k) => k.id === id)
+    if (!key || (region !== 'ai' && region !== 'cn')) return false
+    if (key.region === region) return false
+    key.region = region
+    await this._save()
+    return true
+  }
+
+  async add(rawKey, label = '', region = '') {
+    const key = String(rawKey ?? '').trim()
+    if (key.length < 16) return { ok: false, error: 'Key 太短，看起来不像魔搭令牌' }
+    if (this.has(key)) return { ok: false, error: '这个 Key 已经在池里了' }
+    let id = `k${this.keys.length + 1}`
+    let n = this.keys.length + 1
+    while (this.keys.some((k) => k.id === id)) id = `k${++n}`
+    this.keys.push({
+      id,
+      key,
+      label: String(label ?? '').slice(0, 24),
+      region: region === 'ai' || region === 'cn' ? region : '',
+      addedAt: Date.now(),
+    })
+    this.states.set(id, { until: 0, fails: 0, lastReason: '', lastError: '', used: 0, ok: 0 })
+    await this._save()
+    this._notify()
+    return { ok: true, id }
+  }
+
+  async remove(id) {
+    const before = this.keys.length
+    this.keys = this.keys.filter((k) => k.id !== id)
+    this.states.delete(id)
+    void this._saveCool()
+    if (this.keys.length === before) return false
+    await this._save()
+    this._notify()
+    return true
+  }
+
+  /** 让某个 Key 立刻恢复可用（额度查明后又涨回去了、或手动点了恢复）。 */
+  restore(id) {
+    const st = this.states.get(id)
+    if (!st) return false
+    const was = st.until > 0 || Boolean(st.lastReason)
+    st.until = 0
+    st.lastReason = ''
+    st.lastError = ''
+    st.manual = false
+    if (was) {
+      void this._saveCool()
+      this._notify()
+    }
+    return was
+  }
+
+  async restoreAll() {
+    for (const st of this.states.values()) {
+      st.until = 0
+      st.lastReason = ''
+      st.manual = false
+    }
+    void this._saveCool()
+    this._notify()
+  }
+
+  _notify() {
+    try {
+      this.onChange?.()
+    } catch {}
+  }
+
+  /** 到期自动恢复：任何一次取状态或取 Key 都会顺手清理。 */
+  _expire() {
+    let cooling = 0
+    const now = Date.now()
+    for (const st of this.states.values()) {
+      if (st.until && st.until <= now) {
+        st.until = 0
+        st.lastReason = ''
+      }
+      if (st.until > now) cooling++
+    }
+    return cooling
+  }
+
+  available() {
+    this._expire()
+    const now = Date.now()
+    return this.keys.filter((k) => !(this.states.get(k.id)?.until > now))
+  }
+
+  /** 轮询取下一个可用 Key；exclude 用于同一次请求里不重复试同一个。 */
+  pick(exclude) {
+    const pool = this.available().filter((k) => !exclude?.has(k.id))
+    if (pool.length === 0) return null
+    const key = pool[this.cursor % pool.length]
+    this.cursor = (this.cursor + 1) % Math.max(1, pool.length)
+    return key
+  }
+
+  markSuccess(id) {
+    const st = this.states.get(id)
+    if (!st) return
+    st.used++
+    st.ok++
+    const wasCooling = st.until > 0
+    st.until = 0
+    st.lastReason = ''
+    st.lastError = ''
+    if (wasCooling) void this._saveCool()
+  }
+
+  /**
+   * 停用一个 Key。untilAt 给了就按它恢复（额度用尽 → 次日 00:01），
+   * 否则退回常规冷却。
+   */
+  /**
+   * 手动禁用：一直停用到点了「恢复」为止。
+   * 用 far-until + manual 标记，额度巡检也不会把它自动放开。
+   */
+  disable(id, reason = '手动禁用') {
+    const st = this.states.get(id)
+    if (!st) return false
+    st.until = Date.now() + 100 * 365 * 24 * 3600_000
+    st.lastReason = reason
+    st.lastError = ''
+    st.manual = true
+    void this._saveCool()
+    this._notify()
+    return true
+  }
+
+  /** 手动停用：按给定时间自动恢复（额度类 → 次日 00:01）。 */
+  async stop(id, reason = '手动停用', untilAt) {
+    const st = this.states.get(id)
+    if (!st) return false
+    st.until = Number(untilAt) > Date.now() ? Number(untilAt) : Date.now() + this.cooldownMs
+    st.lastReason = reason
+    st.lastError = ''
+    st.manual = false
+    void this._saveCool()
+    this._notify()
+    return true
+  }
+
+  markFailure(id, reason, error = '', untilAt) {
+    const st = this.states.get(id)
+    if (!st) return 0
+    const now = Date.now()
+    const explicit = Number(untilAt) || 0
+    st.fails++
+    st.used++
+    st.until = explicit > now ? explicit : now + this.cooldownMs
+    st.lastReason = reason
+    st.lastError = String(error ?? '').slice(0, 120)
+    void this._saveCool()
+    this._notify()
+    return st.until
+  }
+
+  status() {
+    const cooling = this._expire()
+    const now = Date.now()
+    const items = this.keys.map((k) => {
+      const st = this.states.get(k.id) ?? {}
+      const left = st.until > now ? st.until - now : 0
+      return {
+        id: k.id,
+        masked: maskKey(k.key),
+        label: k.label,
+        region: k.region ?? '',
+        state: left > 0 ? 'cooling' : 'ok',
+        leftMs: left,
+        untilAt: left > 0 ? st.until : 0,
+        fails: st.fails ?? 0,
+        used: st.used ?? 0,
+        reason: st.lastReason ?? '',
+        lastError: st.lastError ?? '',
+        manual: st.manual === true,
+      }
+    })
+    return {
+      total: items.length,
+      available: items.filter((k) => k.state === 'ok').length,
+      cooling,
+      keys: items,
+      stats: { ...this.stats },
+    }
+  }
+
+  /** 后台心跳：冷却到点自动恢复；有人手改 keys.json 也会被读进来。 */
+  startWatch(getRefresh) {
+    let lastMtime = 0
+    const checkFile = async () => {
+      try {
+        const info = await stat(this.keysFile)
+        const mtime = Number(info.mtimeMs) || 0
+        if (mtime && mtime !== lastMtime) {
+          lastMtime = mtime
+          await this.reload()
+          getRefresh?.()
+        }
+      } catch {}
+    }
+    const fileTimer = setInterval(checkFile, 4000)
+    const coolTimer = setInterval(() => {
+      const cooling = this._expire()
+      const changed = cooling !== this._lastCooldownCount
+      this._lastCooldownCount = cooling
+      if (changed) getRefresh?.()
+    }, 20_000)
+    if (typeof fileTimer.unref === 'function') fileTimer.unref()
+    if (typeof coolTimer.unref === 'function') coolTimer.unref()
+    return () => {
+      clearInterval(fileTimer)
+      clearInterval(coolTimer)
+    }
+  }
+}

package/lib/provider.js

@@ -0,0 +1,227 @@
+// 把本地聚合网关注册成 DSH 的模型供应商：
+// 「模型」页里直接出现「魔搭 Key 池」，选中即可对话，请求仍由网关做多 Key 轮换。
+// 做法与 dsh-connect-modelscope 相同：pi-ai provider + PiAiAdapter。
+import { createRequire } from 'node:module'
+import { pathToFileURL } from 'node:url'
+import { join } from 'node:path'
+
+const HOST_ROOTS = ['/opt/dsh', process.cwd()]
+const DEFAULT_CONTEXT_WINDOW = 131072
+const IMAGE_BUDGETS = {
+  maxRequestImageBytes: 20971520,
+  requestImagePixelBudget: 4194304,
+  requestImageMaxBytes: 1048576,
+}
+
+// 这些类型的模型不能用来对话，列出来只会干扰选择。
+const NON_CHAT = [/(^|[/-])(wan|hunyuanvideo|i2v|t2v|vace|animate|cogvideox)/i, /(image|flux|kolors|sd[-_]?xl|qwen-image|seedream)/i, /(speech|tts|asr|whisper|cosyvoice|audio)/i, /(embedding|rerank|bge-|gte-)/i]
+
+export const isChatModel = (id) => !NON_CHAT.some((re) => re.test(String(id)))
+
+const prettyName = (id) => {
+  const tail = id.includes('/') ? id.slice(id.lastIndexOf('/') + 1) : id
+  return tail.replace(/[-_]/g, ' ').replace(/\s+/g, ' ').trim()
+}
+
+/** 插件不在 profile 的 node_modules 里，宿主包要按 dsh 安装目录兜底解析。 */
+const HOST_SPECIFIER = {
+  '@earendil-works/pi-ai': (root) => `${root}/node_modules/@earendil-works/pi-ai/dist/index.js`,
+  '@earendil-works/pi-ai/api/openai-completions.lazy': (root) =>
+    `${root}/node_modules/@earendil-works/pi-ai/dist/api/openai-completions.lazy.js`,
+  '@deepseek-ai/dsh-llm': (root) => `${root}/node_modules/@deepseek-ai/dsh-llm/lib/index.js`,
+  '@deepseek-ai/dsh-llm-pi-ai': (root) => `${root}/node_modules/@deepseek-ai/dsh-llm-pi-ai/lib/index.js`,
+}
+
+async function loadHost(specifier) {
+  try {
+    return await import(specifier)
+  } catch (first) {
+    for (const root of HOST_ROOTS) {
+      const guess = HOST_SPECIFIER[specifier]?.(root)
+      if (!guess) continue
+      try {
+        return await import(pathToFileURL(guess).href)
+      } catch {}
+      try {
+        const require = createRequire(join(root, 'noop.cjs'))
+        return await import(pathToFileURL(require.resolve(specifier)).href)
+      } catch {}
+    }
+    throw first
+  }
+}
+
+const INERT_AUTH = {
+  credentials: {
+    async read() {},
+    async list() {
+      return []
+    },
+    async modify() {
+      throw new Error('dsh-modelscope-pool has no pi-ai credential lifecycle')
+    },
+    async delete() {},
+  },
+  authContext: { async env() {}, async fileExists() { return false } },
+}
+
+/**
+ * @param ctx 注入过 llm 的上下文
+ * @param options.providerId 供应商 id（出现在模型页）
+ * @param options.displayName 供应商显示名
+ * @param options.getPort 当前网关端口
+ * @param options.getKey 网关口令（没设就是空）
+ * @param options.getModels 当前可用模型 id 列表
+ * @param options.getEntries 带版本/短名的模型条目（优先于 getModels）
+ * @param options.entryId 本插件在 loader 里的条目 id
+ */
+const REGION_TAG = { ai: '国际版', cn: '国内版' }
+
+/** 模型页里显示的名字：短名 + 版本，例如「DeepSeek-V4-Flash-1017 · 国际版」。 */
+const displayNameOf = (entry) => {
+  const base = String(entry?.name ?? '').trim() || prettyName(entry?.id)
+  const regions = Array.isArray(entry?.regions) ? entry.regions : []
+  // 两个站点都有就不加后缀，避免模型页里一堆「· 双站」。
+  if (regions.length === 0 || regions.length > 1) return base
+  const tag = REGION_TAG[regions[0]] ?? ''
+  return tag ? `${base} · ${tag}` : base
+}
+
+export async function registerPoolProvider(ctx, { providerId, displayName, getPort, getKey, getModels, getEntries, entryId, logger }) {
+  if (!ctx.llm || typeof ctx.llm.registerAdapter !== 'function') return false
+  globalThis.__poolProviderDeclared = false
+  // 配置里声明的路由由 pi-ai 在启动早期挂上；我们稍后再决定是否补一个适配器。
+  try {
+    await new Promise((resolve) => setTimeout(resolve, 1200))
+    if ((ctx.llm.listProviders?.() ?? []).some((route) => route.id === providerId)) {
+      globalThis.__poolProviderDeclared = true
+      return 'declared'
+    }
+  } catch {}
+
+  // 配置里已经声明了同名供应商（比如写进 cordis.patch.yml 的同一条），
+  // 宿主会自己把它挂上，重复注册会撞 DUPLICATE_DIRECTORY，直接让位。
+  try {
+    const rows = ctx.llm.listConfigurableProviders?.() ?? []
+    if (rows.some((row) => row.provider === providerId && row.declared === true)) {
+      globalThis.__poolProviderDeclared = true
+      return 'declared'
+    }
+  } catch {}
+
+  let createProvider, openAICompletionsApi, resolveRetryPolicy, PiAiAdapter
+  try {
+    ;[{ createProvider }, { openAICompletionsApi }, { resolveRetryPolicy }, { PiAiAdapter }] = await Promise.all([
+      loadHost('@earendil-works/pi-ai'),
+      loadHost('@earendil-works/pi-ai/api/openai-completions.lazy'),
+      loadHost('@deepseek-ai/dsh-llm'),
+      loadHost('@deepseek-ai/dsh-llm-pi-ai'),
+    ])
+  } catch (error) {
+    globalThis.__poolProviderError = String(error?.message ?? error)
+    logger?.warn?.(`modelscope-pool: 模型供应商注册跳过（${String(error?.message ?? error)}）`)
+    return false
+  }
+
+  const baseUrl = () => `http://127.0.0.1:${getPort() || 8899}/v1`
+  const token = () => String(getKey?.() ?? '').trim() || 'pool-local'
+
+  const entriesOf = () => {
+    const rows = typeof getEntries === 'function' ? getEntries() : []
+    if (Array.isArray(rows) && rows.length > 0) return rows
+    return (typeof getModels === 'function' ? getModels() : []).map((id) => ({ id }))
+  }
+
+  const buildModels = () =>
+    entriesOf()
+      .filter((entry) => isChatModel(entry.id))
+      .map((entry) => ({
+        id: entry.id,
+        name: displayNameOf(entry),
+        api: 'openai-completions',
+        provider: providerId,
+        baseUrl: baseUrl(),
+        input: ['text'],
+        // DSH 拒绝没有正整数上下文的模型，这里给保守值。
+        contextWindow: DEFAULT_CONTEXT_WINDOW,
+        cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
+        reasoning: false,
+      }))
+
+  const provider = {
+    ...createProvider({
+      id: providerId,
+      name: displayName,
+      auth: {
+        apiKey: {
+          name: '本地网关口令',
+          async resolve() {
+            return { auth: { apiKey: token() }, source: displayName }
+          },
+        },
+      },
+      models: buildModels(),
+      api: openAICompletionsApi(),
+    }),
+    getModels: () => buildModels(),
+  }
+
+  const profile = {
+    provider: providerId,
+    displayName,
+    streamIdleTimeoutMs: 180000,
+    retryPolicy: resolveRetryPolicy(undefined, 'dsh-modelscope-pool retryPolicy'),
+    configuredMaxTokens: new Map(),
+    modelErrors: new Map(),
+    defaultContextWindow: DEFAULT_CONTEXT_WINDOW,
+    ...IMAGE_BUDGETS,
+    piProvider: provider,
+  }
+
+  const adapter = new PiAiAdapter({
+    profiles: () => new Map([[providerId, profile]]),
+    auth: INERT_AUTH,
+    resolveApiKey: async () => token(),
+  })
+
+  ctx.effect(() => ctx.llm.registerAdapter([providerId], adapter))
+
+  ctx.effect(() => {
+    if (typeof ctx.llm.registerConfigurableProviders !== 'function') return
+    let off
+    try {
+      off = ctx.llm.registerConfigurableProviders([
+        { provider: providerId, displayName, settingsNs: entryId, settingsPath: [], declared: false },
+      ])
+    } catch (error) {
+      globalThis.__poolProviderError = `directory: ${String(error?.message ?? error)}`
+      logger?.warn?.(`modelscope-pool: 供应商目录注册失败（${String(error?.message ?? error)}）`)
+      return
+    }
+    return () => {
+      try {
+        off?.()
+      } catch {}
+    }
+  })
+
+  ctx.effect(() => {
+    if (typeof ctx.llm.registerModelDiscovery !== 'function') return
+    let off
+    try {
+      off = ctx.llm.registerModelDiscovery(entryId, async () =>
+        buildModels().map((model) => ({ id: model.id, name: model.name, contextWindow: model.contextWindow, inputModalities: ['text'] })),
+      )
+    } catch (error) {
+      logger?.warn?.(`modelscope-pool: 模型发现注册失败（${String(error?.message ?? error)}）`)
+      return
+    }
+    return () => {
+      try {
+        off?.()
+      } catch {}
+    }
+  })
+
+  return true
+}