cf-worker/src/index.js

/**
 * Cloudflare Worker — ds2api 代理健康检查 + HF Space 卡死自愈
 *
 * 部署步骤:
 *   1. cd cf-worker && npm install
 *   2. 设置 secrets:
 *        wrangler secret put DS2API_BASE_URL
 *        wrangler secret put ADMIN_KEY
 *        wrangler secret put VERCEL_BYPASS_TOKEN   (如果开了 Vercel Deployment Protection)
 *        wrangler secret put HF_TOKEN              (HF API token，用于 Space 卡死自愈)
 *   3. wrangler deploy
 *
 * 两个 cron（见 wrangler.toml）：
 *   - 5 小时一次（"0 0,5,10,15,20 * * *" 等价写法）：代理健康检查（runCheck）
 *   - 10 分钟一次：HF Space 卡死检测+自动 restart（runSpaceCheck）
 *
 * 手动触发：
 *   GET/POST https://<worker>/check          → 代理健康检查
 *   GET/POST https://<worker>/check-space    → Space 卡死检测+自愈
 */

const LOGIN_PATH = "/admin/login"
const CHECK_ALL_PATH = "/admin/proxies/check-all"

// Space 卡死自愈相关常量。namespace/name 通过 env.HF_SPACE_REPO
// （格式 "a3216/ds2api"）配置；不配则用默认值。
const DEFAULT_HF_SPACE_REPO = "a3216/ds2api"
// 重启冷却：避免连续 restart 把 Space 打爆。重启后至少等 5 分钟才允许再次重启。
const RESTART_COOLDOWN_MS = 5 * 60 * 1000
// KV 中存上次重启时间的 key。
const LAST_RESTART_KV_KEY = "last_restart"

function vercelHeaders(env) {
  const headers = {}
  const bypass = env.VERCEL_BYPASS_TOKEN
  if (bypass) {
    headers["x-vercel-protection-bypass"] = bypass
  }
  return headers
}

async function login(baseURL, adminKey, extraHeaders) {
  const res = await fetch(`${baseURL}${LOGIN_PATH}`, {
    method: "POST",
    headers: { "Content-Type": "application/json", ...extraHeaders },
    body: JSON.stringify({ admin_key: adminKey, expire_hours: 1 }),
  })
  if (!res.ok) {
    const text = await res.text()
    throw new Error(`login failed (${res.status}): ${text}`)
  }
  const data = await res.json()
  return data.token
}

async function checkAll(baseURL, token, extraHeaders) {
  const res = await fetch(`${baseURL}${CHECK_ALL_PATH}`, {
    method: "POST",
    headers: {
      "Content-Type": "application/json",
      Authorization: `Bearer ${token}`,
      ...extraHeaders,
    },
  })
  if (!res.ok) {
    const text = await res.text()
    throw new Error(`check-all failed (${res.status}): ${text}`)
  }
  return res.json()
}

async function runCheck(env) {
  const baseURL = (env.DS2API_BASE_URL || "").replace(/\/+$/, "")
  const adminKey = env.ADMIN_KEY || ""

  if (!baseURL || !adminKey) {
    throw new Error("DS2API_BASE_URL and ADMIN_KEY secrets must be set")
  }

  const extraHeaders = vercelHeaders(env)
  const token = await login(baseURL, adminKey, extraHeaders)
  const result = await checkAll(baseURL, token, extraHeaders)
  return result
}

// === HF Space 卡死自愈 ===

// 查询 Space runtime 状态。返回 { stage, hardwareCurrent }。
async function fetchSpaceRuntime(repo, token) {
  const url = `https://huggingface.co/api/spaces/${repo}/runtime`
  const res = await fetch(url, {
    headers: { Authorization: `Bearer ${token}` },
  })
  if (!res.ok) {
    const text = await res.text()
    throw new Error(`runtime query failed (${res.status}): ${text}`)
  }
  const data = await res.json()
  return {
    stage: data.stage,
    hardwareCurrent: data.hardware ? data.hardware.current : null,
  }
}

// 判定 Space 是否卡死：非 RUNNING 或未分配硬件。
function isSpaceStuck(runtime) {
  if (!runtime) return true
  if (runtime.stage !== "RUNNING") return true
  if (!runtime.hardwareCurrent) return true
  return false
}

// 触发 HF Space restart。返回 true 表示请求已被接受。
async function restartSpace(repo, token) {
  const url = `https://huggingface.co/api/spaces/${repo}/restart`
  const res = await fetch(url, {
    method: "POST",
    headers: { Authorization: `Bearer ${token}` },
  })
  if (!res.ok) {
    const text = await res.text()
    throw new Error(`restart failed (${res.status}): ${text}`)
  }
  return true
}

// 从 KV 读上次重启时间戳（ms）。KV 未绑定或无值返回 0。
async function getLastRestart(kv) {
  if (!kv) return 0
  try {
    const v = await kv.get(LAST_RESTART_KV_KEY)
    const n = parseInt(v, 10)
    return Number.isFinite(n) ? n : 0
  } catch {
    return 0
  }
}

async function setLastRestart(kv) {
  if (!kv) return
  try {
    await kv.put(LAST_RESTART_KV_KEY, String(Date.now()))
  } catch {
    // KV 写失败不阻断主流程。
  }
}

// Space 卡死检测+自愈主逻辑。
async function runSpaceCheck(env) {
  const token = env.HF_TOKEN || ""
  if (!token) {
    throw new Error("HF_TOKEN secret must be set for space self-healing")
  }
  const repo = (env.HF_SPACE_REPO || DEFAULT_HF_SPACE_REPO).trim()

  const runtime = await fetchSpaceRuntime(repo, token)
  if (!isSpaceStuck(runtime)) {
    return {
      ok: true,
      action: "none",
      repo,
      stage: runtime.stage,
      hardwareCurrent: runtime.hardwareCurrent,
      message: "space healthy, no action",
    }
  }

  // 卡死，检查冷却。
  const last = await getLastRestart(env.SPACE_STATE)
  const elapsed = Date.now() - last
  if (elapsed < RESTART_COOLDOWN_MS) {
    const waitMs = RESTART_COOLDOWN_MS - elapsed
    return {
      ok: true,
      action: "cooldown",
      repo,
      stage: runtime.stage,
      hardwareCurrent: runtime.hardwareCurrent,
      message: `space stuck but in restart cooldown, wait ${Math.ceil(waitMs / 1000)}s`,
    }
  }

  // 触发重启。
  await restartSpace(repo, token)
  await setLastRestart(env.SPACE_STATE)
  return {
    ok: true,
    action: "restarted",
    repo,
    stage: runtime.stage,
    hardwareCurrent: runtime.hardwareCurrent,
    message: "space was stuck, restart triggered",
  }
}

// 区分两种 cron：5 小时一次的代理检查 vs 10 分钟一次的 Space 检查。
// scheduled.cron 即触发的 cron 表达式（字符串）。
function isSpaceCheckCron(cron) {
  return cron === "*/10 * * * *"
}

export default {
  async scheduled(event, env, ctx) {
    if (isSpaceCheckCron(event.cron)) {
      ctx.waitUntil(
        runSpaceCheck(env)
          .then((r) => {
            if (r.action === "restarted") {
              console.warn(`[space] ${r.message} repo=${r.repo} stage=${r.stage} hardware=${r.hardwareCurrent}`)
            } else {
              console.log(`[space] ${r.message} repo=${r.repo} stage=${r.stage} hardware=${r.hardwareCurrent}`)
            }
          })
          .catch((err) => {
            console.error(`[space] check failed: ${err.message}`)
          })
      )
      return
    }

    // 默认：代理健康检查 cron。
    ctx.waitUntil(
      runCheck(env)
        .then((result) => {
          const items = result.items || []
          const healthy = items.filter((i) => i.healthy && !i.disabled).length
          const banned = items.filter((i) => i.disabled).length
          console.log(
            `proxy health check done: ${healthy}/${items.length} healthy, ${banned} banned`
          )
        })
        .catch((err) => {
          console.error(`proxy health check failed: ${err.message}`)
        })
    )
  },

  async fetch(request, env) {
    const url = new URL(request.url)

    if (url.pathname === "/check") {
      try {
        const result = await runCheck(env)
        return Response.json(result)
      } catch (err) {
        return Response.json({ error: err.message }, { status: 500 })
      }
    }

    if (url.pathname === "/check-space") {
      try {
        const result = await runSpaceCheck(env)
        return Response.json(result)
      } catch (err) {
        return Response.json({ error: err.message }, { status: 500 })
      }
    }

    return new Response("not found", { status: 404 })
  },
}