dsh-account-pool 0.1.0

package/lib/shim.js

@@ -0,0 +1,370 @@
+/**
+ * Loopback shim：插件自带的「网关」。
+ *
+ * 为什么需要它：
+ *   pi-ai provider 只认「一个 baseURL + 一个 apiKey」。但我们要在多账号之间
+ *   自动切换，而切换必须发生在每次请求——不能让 pi-ai 知道账号的存在。
+ *   所以在 127.0.0.1 上起一个 OpenAI 兼容端点，pi-ai 指向它；
+ *   它收到请求后自己选号、转发上游、失败就换号重试。
+ *
+ *   DSH → pi-ai adapter → 本 shim（127.0.0.1:随机端口）→ copilot.tencent.com
+ *                          ↑ 自动切换就发生在这里
+ *
+ * 安全加固（参照 dsh-connect-workbuddy，安全代码不做「改善」）：
+ *   - 只绑 127.0.0.1，绝不监听其他网卡
+ *   - 随机端口，随机进程内 secret 作 Bearer；真 token 不交给 pi-ai
+ *   - Host 必须回环（防 DNS rebinding）
+ *   - Origin 若存在必须回环（防浏览器跨站请求）
+ *   - 常量时间比较 secret（防时序侧信道）
+ *   - 请求体上限，防内存打爆
+ */
+
+import { randomBytes, timingSafeEqual } from 'node:crypto'
+import { createServer } from 'node:http'
+
+import {
+  chatStream as workbuddyChatStream,
+  prepareChatBody as workbuddyPrepareBody,
+  isRetryableWithAnotherAccount as workbuddyRetryable,
+} from './upstream.js'
+import { extractUsage } from './usage.js'
+
+/** 请求体上限 64 MiB（与参照项目一致，长上下文请求会比较大）。 */
+const REQUEST_BODY_LIMIT = 64 * 1024 * 1024
+
+/** 回环主机名白名单。 */
+const LOOPBACK_HOSTS = new Set(['127.0.0.1', 'localhost', '[::1]'])
+
+/** 最多为一次请求尝试几个账号。防止全挂时无限重试。 */
+const MAX_ACCOUNT_ATTEMPTS = 3
+
+/** 上游失败类型 → 回给 pi-ai 的 HTTP 状态码。 */
+const KIND_STATUS = {
+  hard_credit: 402,
+  soft_rate: 429,
+  session_dead: 401,
+  not_found: 502,
+  server: 502,
+  client: 400,
+}
+
+/** 从 Host 头里剥掉端口，IPv6 方括号要特殊处理。 */
+function hostnameOfHost(host) {
+  let value = String(host).trim().toLowerCase()
+  if (value.startsWith('[')) {
+    const end = value.indexOf(']')
+    return end === -1 ? value : value.slice(0, end + 1)
+  }
+  const colon = value.lastIndexOf(':')
+  if (colon !== -1 && /^\d+$/.test(value.slice(colon + 1))) value = value.slice(0, colon)
+  return value
+}
+
+/** Host 必须指向回环——DNS rebinding 攻击会带自己的域名，这里挡掉。 */
+function hostIsLoopback(host) {
+  if (host === undefined || String(host).trim() === '') return false
+  return LOOPBACK_HOSTS.has(hostnameOfHost(host))
+}
+
+/** 浏览器发的 Origin 必须在回环；插件自己的 fetch 不带 Origin，放行。 */
+function originIsLoopback(origin) {
+  if (origin === undefined || String(origin).trim() === '') return true
+  try {
+    const { hostname } = new URL(origin)
+    return LOOPBACK_HOSTS.has(hostname) || hostname === '::1'
+  } catch {
+    return false
+  }
+}
+
+/** 常量时间比较，避免逐字节比较泄露 secret。 */
+function secretMatches(provided, expected) {
+  const a = Buffer.from(String(provided ?? ''))
+  const b = Buffer.from(expected)
+  if (a.length !== b.length) return false
+  return timingSafeEqual(a, b)
+}
+
+/** 从 Authorization 头取出 Bearer token。 */
+function bearerOf(headers) {
+  const raw = headers.authorization
+  if (typeof raw !== 'string') return ''
+  const match = /^Bearer\s+(.+)$/i.exec(raw.trim())
+  return match === null ? '' : match[1]
+}
+
+function writeJson(res, status, body) {
+  const payload = JSON.stringify(body)
+  res.writeHead(status, {
+    'content-type': 'application/json',
+    'content-length': Buffer.byteLength(payload),
+  })
+  res.end(payload)
+}
+
+/** 按 OpenAI 错误格式回错，pi-ai 才能正确分类。 */
+function writeOpenAIError(res, status, kind, message) {
+  writeJson(res, status, { error: { message, type: kind, code: kind } })
+}
+
+/** 读取并限制请求体大小。 */
+async function readBody(req) {
+  const chunks = []
+  let size = 0
+  for await (const chunk of req) {
+    size += chunk.length
+    if (size > REQUEST_BODY_LIMIT) throw new Error('请求体超过上限')
+    chunks.push(chunk)
+  }
+  return Buffer.concat(chunks).toString('utf8')
+}
+
+/**
+ * 起一个 shim。
+ *
+ * @param {object} options
+ * @param {() => Promise<Array<{id:string, credential:object}>>} options.listAccounts
+ *        返回本次可选的账号（已含可用 token）。每次请求现取，保证拿到最新 token。
+ * @param {import('./selector.js').AccountPicker} options.picker 选号器
+ * @param {{info:Function, warn:Function, error:Function}} options.logger
+ * @param {object} [options.usage] 用量记录器（有就记，没有就跳过）
+ * @param {object} [options.adapter] 上游适配器；不传则用 WorkBuddy 的默认实现。
+ *   三个成员：
+ *     chatStream(credential, bodyJson, signal) → 上游响应
+ *     prepareBody(rawBody) → 改写成该上游认的请求体
+ *     retryable(kind) → 该失败是否值得换号重试
+ * @returns {Promise<{baseUrl:()=>string, token:()=>string, close:()=>Promise<void>}>}
+ */
+export async function createShim({ listAccounts, picker, logger, usage, adapter }) {
+  // 上游差异全部收敛在这个适配器里——选号、重试、用量、流处理都是共用的。
+  // Trae 与 WorkBuddy 的协议不同（头部、请求体改写、错误分类），
+  // 但「多个账号之间怎么选、失败怎么换」这套逻辑完全相同。
+  const upstream = {
+    chatStream: adapter?.chatStream ?? workbuddyChatStream,
+    prepareBody: adapter?.prepareBody ?? workbuddyPrepareBody,
+    retryable: adapter?.retryable ?? workbuddyRetryable,
+  }
+  /** 进程内随机 secret：pi-ai 拿它当 apiKey，真 token 不出去。 */
+  const secret = randomBytes(32).toString('hex')
+
+  /**
+   * 处理一次 chat 请求。核心：选号 → 转发 → 失败换号重试。
+   */
+  async function handleChat(req, res, rawBody) {
+    const sessionKeyOf = (body) => {
+      // DSH 会把会话信息透传下来，用它做粘性键；没有就退化为不粘。
+      return body?.metadata?.conversation_id
+        ?? body?.metadata?.conversationId
+        ?? body?.conversation_id
+        ?? body?.conversationId
+        ?? ''
+    }
+
+    let body
+    try {
+      body = JSON.parse(rawBody)
+    } catch {
+      writeOpenAIError(res, 400, 'client', '请求体不是合法 JSON')
+      return
+    }
+
+    const sessionKey = sessionKeyOf(body)
+    const requestedModel = typeof body?.model === 'string' ? body.model : '(unknown)'
+    const prepared = upstream.prepareBody(rawBody)
+
+    // 每次请求现取账号，保证用的是最新 token。
+    let accounts
+    try {
+      accounts = await listAccounts()
+    } catch (error) {
+      writeOpenAIError(res, 502, 'server', `无法读取账号：${error?.message ?? error}`)
+      return
+    }
+    if (accounts.length === 0) {
+      writeOpenAIError(res, 401, 'session_dead', '没有已登录的 WorkBuddy 账号，请到插件设置页添加')
+      return
+    }
+
+    // 客户端断开（DSH 里点「停止生成」）时要中止上游请求。
+    // 否则上游流会继续跑完，白白消耗积分。
+    const controller = new AbortController()
+    let clientGone = false
+    res.on('close', () => {
+      // 响应正常结束时也会触发 close，用 writableEnded 区分
+      if (!res.writableEnded) {
+        clientGone = true
+        controller.abort()
+      }
+    })
+
+    const candidates = accounts.map(a => a.id)
+    const tried = new Set()
+    let lastFailure
+
+    // 依次尝试：每次挑一个没试过的号，失败就换下一个。
+    for (let attempt = 0; attempt < Math.min(MAX_ACCOUNT_ATTEMPTS, candidates.length); attempt += 1) {
+      // 客户端已经断开就没必要再换号重试了。
+      if (clientGone) break
+      const remaining = candidates.filter(id => !tried.has(id))
+      if (remaining.length === 0) break
+
+      const accountId = picker.pick(remaining, sessionKey)
+      if (accountId === undefined) break
+      tried.add(accountId)
+      const startedAt = Date.now()
+
+      const account = accounts.find(a => a.id === accountId)
+      if (account === undefined) continue
+
+      const result = await upstream.chatStream(account.credential, prepared, controller.signal)
+
+      if (result.ok) {
+        picker.reportSuccess(accountId)
+        // 首帧前的错误已经排除，直接把上游 SSE 流回给 pi-ai。
+        res.writeHead(200, {
+          'content-type': 'text/event-stream; charset=utf-8',
+          'cache-control': 'no-cache',
+          connection: 'keep-alive',
+        })
+        // 边转发边抓 usage：上游在流末尾发一帧带 usage（含 credit）的数据。
+        // 用 TextDecoder 累积文本，因为 usage 帧可能跨 chunk 边界。
+        const decoder = new TextDecoder()
+        let tail = ''
+        let captured
+        try {
+          for await (const chunk of result.response.body) {
+            res.write(chunk)
+            if (usage !== undefined) {
+              tail += decoder.decode(chunk, { stream: true })
+              const found = extractUsage(tail)
+              if (found !== undefined) captured = found
+              // 控制内存：只在**换行边界**截断。
+              // 直接 slice 会把一个 SSE 帧切成两半，若 usage 帧正好跨在
+              // 切点上就永远抓不到——按行切不会破坏帧的完整性。
+              if (tail.length > 64 * 1024) {
+                const boundary = tail.lastIndexOf('\n')
+                // 保留最后一段完整行；万一整块没有换行（异常），才退化为保留尾部
+                tail = boundary > 0 ? tail.slice(boundary) : tail.slice(-8192)
+              }
+            }
+          }
+        } catch (error) {
+          // 客户端主动断开是正常情况（abort），不当异常报
+          if (!clientGone) logger.warn(`上游流出错：${error?.message ?? error}`)
+        }
+        if (!clientGone) res.end()
+        // 记录用量。ok 以「是否拿到 usage」判定：拿不到说明这次尝试
+        // 没成功产出 token。
+        // 只有成功请求才传 latencyMs——见下面失败分支的说明。
+        if (usage !== undefined) {
+          usage.add({
+            accountId,
+            model: captured?.model ?? requestedModel,
+            ok: captured?.usage !== undefined,
+            usage: captured?.usage,
+            latencyMs: Date.now() - startedAt,
+          })
+        }
+        return
+      }
+
+      // 失败：处置账号，记录原因。
+      const action = picker.reportFailure(accountId, result.kind)
+      logger.warn(`账号 ${accountId} 请求失败（${result.kind}）：${action}`)
+      // 失败的尝试计入请求数——重试放大正是靠这一列才看得出来。
+      //
+      // 但**不记耗时**：失败往往在上游快速返回（429/4xx 只要几百毫秒），
+      // 而成功要跑完整条流（数千毫秒）。混在一起算会把平均延迟拉低，
+      // 按小时分桶时更会忽高忽低，看不出真实响应速度。
+      // 「平均延迟」的语义是「请求要等多久」，失败的等待没有意义。
+      if (usage !== undefined) {
+        usage.add({
+          accountId,
+          model: requestedModel,
+          ok: false,
+        })
+      }
+      lastFailure = result
+
+      // 请求本身有问题的话换号也没用，直接回错，省得白打一遍。
+      if (!upstream.retryable(result.kind)) break
+    }
+
+    if (lastFailure === undefined) {
+      writeOpenAIError(res, 502, 'server', '所有账号都不可用')
+      return
+    }
+    writeOpenAIError(
+      res,
+      KIND_STATUS[lastFailure.kind] ?? 502,
+      lastFailure.kind,
+      `所有账号尝试失败：${lastFailure.message}`,
+    )
+  }
+
+  const server = createServer((req, res) => {
+    void (async () => {
+      try {
+        // ---- 入站加固：四重校验，任一不过直接拒绝 ----
+        if (!hostIsLoopback(req.headers.host)) {
+          writeJson(res, 403, { error: { message: 'Host 非回环', type: 'forbidden', code: 'forbidden' } })
+          return
+        }
+        if (!originIsLoopback(req.headers.origin)) {
+          writeJson(res, 403, { error: { message: 'Origin 非回环', type: 'forbidden', code: 'forbidden' } })
+          return
+        }
+        const url = new URL(req.url ?? '/', 'http://127.0.0.1')
+        if (!url.pathname.endsWith('/chat/completions')) {
+          writeJson(res, 404, { error: { message: '未知路径', type: 'not_found', code: 'not_found' } })
+          return
+        }
+        if (req.method !== 'POST') {
+          res.writeHead(405, { allow: 'POST' })
+          res.end()
+          return
+        }
+        if (!secretMatches(bearerOf(req.headers), secret)) {
+          writeJson(res, 401, { error: { message: 'secret 不匹配', type: 'unauthorized', code: 'unauthorized' } })
+          return
+        }
+        const contentType = String(req.headers['content-type'] ?? '')
+        if (!contentType.toLowerCase().startsWith('application/json')) {
+          writeJson(res, 415, { error: { message: '需要 application/json', type: 'unsupported', code: 'unsupported' } })
+          return
+        }
+
+        const rawBody = await readBody(req)
+        await handleChat(req, res, rawBody)
+      } catch (error) {
+        logger.error(`shim 处理请求异常：${error?.stack ?? error}`)
+        if (!res.headersSent) {
+          writeJson(res, 500, { error: { message: '内部错误', type: 'server', code: 'server' } })
+        } else {
+          res.end()
+        }
+      }
+    })()
+  })
+
+  // 绑随机端口、只绑回环。
+  await new Promise((resolve, reject) => {
+    server.once('error', reject)
+    server.listen(0, '127.0.0.1', () => {
+      server.removeListener('error', reject)
+      resolve()
+    })
+  })
+
+  const address = server.address()
+  const origin = `http://127.0.0.1:${address.port}`
+
+  return {
+    baseUrl: () => origin,
+    token: () => secret,
+    close: () => new Promise((resolve) => {
+      server.closeAllConnections?.()
+      server.close(() => resolve())
+    }),
+  }
+}