宜博 / 2080tiserver

#!/usr/bin/env python3
"""H3 Console：给 gateway.py 加一层 Web UI + 上传 + 任务持久化。

不重做调度——gateway.py 已经做了选后端/分区路由/限流，这里只反代它的
/v1/videos* 接口，加三样它没有的东西：文件上传（返回可喂给 api.py 的 URL）、
SQLite 任务影子记录（gateway.py 的 JOBS 在后端进程内存里，进程重启就丢）、
静态前端。

弹性 partition 调度器（kill+restart 某个 GPU 后端切 fl2va/ref2va）只在
H3_ENABLE_SUPERVISOR=1 时启用——这一层直接操作服务器上的真实进程，本地
Mac 开发环境没有意义，见 DESIGN.md 「四点五」。
"""
import os, re, sqlite3, time, uuid, json, asyncio, subprocess
from pathlib import Path
from typing import Optional

import hmac
import httpx
from fastapi import APIRouter, Depends, FastAPI, HTTPException, UploadFile, File, Request
from fastapi.responses import StreamingResponse, FileResponse, HTMLResponse
from fastapi.staticfiles import StaticFiles
from PIL import Image

BASE = Path(__file__).resolve().parent
ASSETS = BASE / "assets"
ASSETS.mkdir(exist_ok=True)
DB_PATH = BASE / "console.sqlite3"

GATEWAY_URL = os.environ.get("H3_GATEWAY_URL", "http://192.168.110.67:8000")
CONSOLE_PUBLIC_BASE = os.environ.get("H3_CONSOLE_PUBLIC_BASE")  # 供后端 worker 回读上传图片用；见 README
MAX_UPLOAD_BYTES = 20 * 1024 * 1024  # 20MB，够单张参考图/首尾帧，防止占满磁盘
MAX_VIDEO_UPLOAD_BYTES = 100 * 1024 * 1024  # 100MB，参考视频（2-15秒）比图片大得多
MAX_AUDIO_UPLOAD_BYTES = 30 * 1024 * 1024  # 30MB，独立音频参考
ENABLE_SUPERVISOR = os.environ.get("H3_ENABLE_SUPERVISOR") == "1"
# 部署在共享域名的子路径下时设这个（比如挂在 qwen.wudaima.com/console 而不是独立域名），
# 见 README「子路径部署」。留空就是原来的行为，挂在根路径，本地/局域网直连不受影响。
BASE_PATH = os.environ.get("H3_CONSOLE_BASE_PATH", "").rstrip("/")
# 公网直接可达（qwen.wudaima.com/console 无需登录就能提交任务/看到所有历史任务），设了这个
# 才校验 /v1/* 接口；留空等于原来的行为（局域网/本地开发不用配）。`<video src>`/下载链接是
# 浏览器原生请求，加不上 Authorization 头，所以顺带认 `?token=` 查询参数——跟 header 校验
# 走同一份比较逻辑，不是单独开了个后门。
AUTH_TOKEN = os.environ.get("H3_CONSOLE_AUTH_TOKEN", "")


def _require_token(request: Request):
    if not AUTH_TOKEN:
        return
    got = request.query_params.get("token", "")
    auth = request.headers.get("authorization", "")
    if auth.startswith("Bearer "):
        got = auth[len("Bearer "):]
    if not hmac.compare_digest(got, AUTH_TOKEN):
        raise HTTPException(401, "缺少或错误的访问令牌（Authorization: Bearer <token> 或 ?token=）")

# GPU/端口映射，对应服务器上 start-all.sh 写死的分配。仅 supervisor 用得到。
BACKEND_TOPOLOGY = [
    {"gpu": 1, "port": 8010},
    {"gpu": 2, "port": 8011},
    {"gpu": 0, "port": 8012},
    {"gpu": 3, "port": 8013},
]
H3_ROOT = Path("/data/models/minimax-h3")  # 仅服务器上存在，本地开发不会用到

app = FastAPI(title="H3 Console")
client = httpx.AsyncClient(timeout=30.0)
# 2026-08-24："重置GPU"/"停止运行中任务"都是硬重启整个raylight-server，进程重启期间端口
# 直接连不上，这段时间提交新任务只会拿到一个费解的连接失败——干脆记一个"重置中"窗口，
# 提交请求直接给个说得清楚的409，不留给用户猜。25秒是给硬重启(~15-20秒)留够余量的估算值，
# 不是精确等到端口真正恢复健康。
RESET_WINDOW_SECONDS = 25
_reset_until = 0.0


def _mark_resetting():
    global _reset_until
    _reset_until = time.time() + RESET_WINDOW_SECONDS


def _check_not_resetting():
    remaining = _reset_until - time.time()
    if remaining > 0:
        raise HTTPException(409, f"GPU正在重置/服务重启中，还剩约{int(remaining)}秒，请稍后重试")
# 所有 /v1/* 接口挂在这个 router 上而不是直接挂在 app 上，是为了能整体套一层 BASE_PATH
# 前缀（app.include_router(router, prefix=BASE_PATH)，见文件末尾）——BASE_PATH 留空时
# 前缀是空字符串，行为和直接挂在 app 上完全一样。
router = APIRouter()  # 页面路由（/、/newtask.html、/task.html/{tid}）——不校验token，
# 不然连页面壳子都加载不出来，无从输入token
api = APIRouter(dependencies=[Depends(_require_token)])  # 所有 /v1/* 接口


def db():
    # 2026-09-05：console 是 8080(LAN)/8081(公网) 两个独立进程共用同一份 sqlite 文件，各自的
    # reconciler_loop 每 5 秒都会对着"进行中的任务"逐条 UPDATE——队列一旦攒到 200+（一次
    # fire-and-forget 式并发提交 150 个任务撞出来的），两个进程的写请求频繁互相撞车。默认
    # journal_mode=DELETE 下同一时刻只允许一个写者、读写还会互相阻塞，且没设 timeout 时
    # sqlite3 的默认等锁时间只有 5 秒，稍微堆一下就直接抛"database is locked"，而不是排队
    # 等一下——GET /v1/backends 这种只读查询也会被前面攒着的写请求堵住，实测卡到 38 秒。
    # WAL 模式允许"多个读者 + 一个写者"并发不互斥，配合更长的 busy timeout 让写者互相排队
    # 重试而不是直接报错，两个改动都是连接级别、幂等，不需要额外迁移。
    con = sqlite3.connect(DB_PATH, timeout=30)
    con.execute("PRAGMA journal_mode=WAL")
    con.row_factory = sqlite3.Row
    return con


def init_db():
    with db() as con:
        con.execute("""
            CREATE TABLE IF NOT EXISTS jobs (
                id TEXT PRIMARY KEY,
                gateway_id TEXT,
                partition TEXT NOT NULL,
                backend TEXT,
                status TEXT NOT NULL,
                req_json TEXT NOT NULL,
                created_at REAL NOT NULL,
                updated_at REAL NOT NULL,
                last_known_json TEXT
            )
        """)
        # 老库升级：CREATE TABLE IF NOT EXISTS 不会给已存在的表加列，只能单独 ALTER。
        # 第二次启动起这行必然报"列已存在"，吞掉就行，不影响其它列。
        try:
            con.execute("ALTER TABLE jobs ADD COLUMN deleted_at REAL")
        except sqlite3.OperationalError:
            pass


init_db()


# ---------------- 资产上传 ----------------
# api.py 的 GenReq 只认 URL 字段（first_frame_url / reference_image_urls），内部用
# urllib.urlopen 拉取。所以上传接口的活是「收文件 -> 落盘 -> 吐一个后端能拉到的 URL」，
# 不改 api.py 一行代码。
IMAGE_EXTS = (".jpg", ".jpeg", ".png", ".webp")
VIDEO_EXTS = (".mp4", ".mov", ".webm", ".mkv")
AUDIO_EXTS = (".mp3", ".wav", ".m4a", ".ogg", ".flac")


@api.post("/v1/assets")
async def upload_asset(file: UploadFile = File(...)):
    ext = Path(file.filename or "").suffix.lower()
    # 2026-08-23：原来只收图片，现在参考视频/独立音频参考也走这个接口——按扩展名分流到
    # 各自的大小上限，视频/音频没有 PIL 那样现成的解码校验，先不做，坏文件会在 ComfyUI
    # LoadVideo/LoadAudio 那一步报错，属于可接受的失败模式。
    if ext in IMAGE_EXTS:
        kind, max_bytes = "image", MAX_UPLOAD_BYTES
    elif ext in VIDEO_EXTS:
        kind, max_bytes = "video", MAX_VIDEO_UPLOAD_BYTES
    elif ext in AUDIO_EXTS:
        kind, max_bytes = "audio", MAX_AUDIO_UPLOAD_BYTES
    else:
        raise HTTPException(400, f"不支持的文件类型 {ext}，只收图片(jpg/png/webp)/"
                                  f"视频(mp4/mov/webm/mkv)/音频(mp3/wav/m4a/ogg/flac)")
    body = await file.read()
    if len(body) > max_bytes:
        raise HTTPException(413, f"文件超过 {max_bytes // 1024 // 1024}MB 上限")
    aid = uuid.uuid4().hex[:16]
    # 原始文件名塞进存盘文件名里（aid__原名.ext），前端从 URL 就能反解出来显示——
    # 不用另外给 job req 加字段存文件名，省了一次 schema 改动。
    stem = re.sub(r"[^\w.-]+", "_", Path(file.filename or "").stem)[:60] or "file"
    path = ASSETS / f"{aid}__{stem}{ext}"
    path.write_bytes(body)
    if kind == "image":
        try:
            Image.open(path).verify()  # 校验确实是能解码的图片，不是随便什么二进制
        except Exception:
            path.unlink(missing_ok=True)
            raise HTTPException(400, "文件不是有效图片")
    return {"id": aid, "url": f"{_public_base()}/assets/{aid}__{stem}{ext}", "kind": kind}


def _public_base() -> str:
    if CONSOLE_PUBLIC_BASE:
        return CONSOLE_PUBLIC_BASE.rstrip("/")
    # 本地开发时没配这个环境变量会直接报错提示——见 README「已知限制」：
    # 上传图片的 URL 必须是 api.py 所在服务器能反向拉取到的地址，本地 Mac 默认猜不出这个地址。
    raise HTTPException(
        500, "未设置 H3_CONSOLE_PUBLIC_BASE：上传接口需要知道自己对外可访问的地址，"
             "这个地址会被服务器上的 api.py worker 反向拉取，见 README。")


app.mount(f"{BASE_PATH}/assets", StaticFiles(directory=str(ASSETS)), name="assets")


# ---------------- 生成任务：console 自己是排队的主人，gateway.py 只是"当前有空的执行者" ----------------
# 关键点：gateway.py 对没有就绪后端的 partition 直接 503（`没有就绪的 fl2va 实例`），不排队。
# 如果 console 只是把这个 503 原样透传给用户，"有几个任务排队就用几张卡跑 fl2va"就无从谈起——
# 请求在 503 那一步就死了，根本没机会被弹性调度器看到、更别说触发 flip。
# 所以 console 自己维护身份（console 生成的 id）和排队状态：提交时先落 sqlite 记一条 status=queued，
# 尝试提交给 gateway；提交不上（503/429）就留在本地队列里，交给下面的 reconciler_loop 每隔几秒重试，
# 弹性调度器看到的 pending 计数正是这些「还没提交上去」的本地排队任务。
def _partition_of(req: dict) -> str:
    # 2026-08-23 修复：原来是二选一（没有first/last frame就当ref2va），会把纯文生视频
    # （reference_image_urls也是空的）误标成"ref2va"——跟前端 rerunJob() 自己的 kind 判断
    # （index.html 里早就正确区分 text/fl2va/ref2va 三种）对不上，在 task.html 上看起来
    # "标着ref2va却没有参考图"，容易让人怀疑参考图是不是丢了/没显示出来，其实只是这个分类
    # 少了纯文生视频这一档。改成跟前端一致的三分类。
    if req.get("first_frame_url") or req.get("last_frame_url"):
        return "fl2va"
    # 2026-08-23：新增参考视频/独立音频参考字段时漏了这里，跟 raylight/server.py 自己的
    # _partition_of 对不上——那边路由图用的是它自己那份判断（功能不受影响，只是这里的
    # 展示标签会显示错），保持两份判断逻辑一致。
    if req.get("reference_image_urls") or req.get("reference_video_urls") or req.get("reference_audio_urls"):
        return "ref2va"
    return "text"


def _row_view(row) -> dict:
    """把一行 sqlite 记录还原成对外的任务视图，用 console 自己的 id，不暴露 gateway 内部 id。"""
    if row["last_known_json"]:
        view = json.loads(row["last_known_json"])
    else:
        view = {}
    view["id"] = row["id"]
    view["status"] = row["status"]
    view["partition"] = row["partition"]
    view["backend"] = row["backend"]
    view.setdefault("req", json.loads(row["req_json"]))
    view["created_at"] = row["created_at"]
    view["deleted_at"] = row["deleted_at"]
    # gateway.py/api.py 原样带过来的 `poll`/`download` 字段用的是 gateway 内部 id（不是 console
    # 自己的 id），照抄会指向一个 console 自己的 sqlite 里根本查不到的路径，直接 404——这两个
    # 字段本来就该跟 `id` 一样，统一改成用 console 的 id 重新拼。
    if "poll" in view:
        view["poll"] = f"{BASE_PATH}/v1/videos/{row['id']}"
    if "download" in view:
        view["download"] = f"{BASE_PATH}/v1/videos/{row['id']}/file"
    return view


# create_video 首次提交、get_video 轮询时顺手提交、reconciler_loop 每 5 秒扫描，三条路都会在
# gateway_id 还是 NULL 的窗口里调 _try_submit(cid)。三个协程共享同一个事件循环，`await
# client.post(...)` 让出控制权的那段时间里，另外两条路完全可能也读到「还没提交过」，
# 于是同一个 cid 被 POST 给 gateway 两次——两张卡各烧一遍 GPU，console 只记得住最后一次
# UPDATE 写回的 gateway_id，另一条产物是真实存在但没人知道、也没人能停的孤儿任务。
# 2026-08-20 实测踩到过：一个 15s/832×480 的探测任务真的在两张卡上各 OOM 了一次。
# 修法：进程内占坑（不是数据库锁——server.py 单进程单事件循环，内存里一个 set 就够，
# 两行判断中间没有 await，天然原子，不会被其它协程插进来）。
_SUBMITTING: set[str] = set()

# 2026-09-04：上面这把 _SUBMITTING 只挡得住"同一个 cid 被提交两次"，挡不住另一种真实踩到过的
# 重复——_try_submit 从发 POST 到把 gateway_id 写回数据库之间有一段 await 窗口，这段时间里
# gateway(raylight) 那边这条任务已经创建、能在它自己的 /v1/videos 列表里查到了，但 console 的
# jobs 表还没记上这个 gateway_id。如果 reconciler_loop 的 _discover_from_gateway() 正好在这个
# 窗口执行，会把这条"看起来没人认领"的 gateway 任务当成孤儿任务，拿 gateway_id 本身当 console id
# 再插一行——于是同一个真实生成任务在任务列表里出现两条记录，状态/进度看起来一样，其实是同一个
# comfy_prompt_id（案例：4c99f52a3a7a / ff818215fd3e，2026-09-04）。
# 修法：把"提交"（发请求→写回 gateway_id）和"发现孤儿"（读已知 gateway_id 集合→插入孤儿行）
# 这两段簿记操作用同一把锁互斥，全程不重叠——反正 raylight 本来就一次只处理一个请求，串行化
# 这两段轻量 HTTP 调用不影响真正的生成吞吐量。
_SUBMIT_LOCK = asyncio.Lock()


async def _try_submit(cid: str):
    """把一条本地排队的任务尝试提交给 gateway。给 create_video 首次调用和 reconciler_loop 复用。"""
    if cid in _SUBMITTING:
        return  # 已经有一次提交在飞行中，等它写完 gateway_id 再说，不要并发再发一次
    with db() as con:
        row = con.execute("SELECT * FROM jobs WHERE id=?", (cid,)).fetchone()
    if row is None or row["gateway_id"] is not None:
        return  # 任务不存在，或已经提交过
    _SUBMITTING.add(cid)
    try:
        async with _SUBMIT_LOCK:
            req = json.loads(row["req_json"])
            try:
                resp = await client.post(f"{GATEWAY_URL}/v1/videos", json=req)
            except httpx.HTTPError:
                return  # gateway 暂时连不上，下一轮再试
            now = time.time()
            if resp.status_code < 400:
                job = resp.json()
                gateway_id = job["id"]
                job["id"] = cid  # 对外只暴露 console 自己的 id
                with db() as con:
                    con.execute(
                        "UPDATE jobs SET gateway_id=?, backend=?, status=?, updated_at=?, last_known_json=? WHERE id=?",
                        (gateway_id, job.get("backend"), job.get("status", "queued"),
                         now, json.dumps(job, ensure_ascii=False), cid),
                    )
            elif resp.status_code in (503, 429):
                return  # 没有就绪后端 / 队列满，留在本地队列，等 flip 或别的任务跑完再试
            else:
                # 真正的客户端错误（比如 422 参数越界、400 turbo+steps>8），重试也不会成功，直接判失败
                err_view = {"id": cid, "status": "failed", "error": resp.text,
                            "req": req, "partition": row["partition"], "created_at": row["created_at"]}
                with db() as con:
                    con.execute("UPDATE jobs SET status='failed', updated_at=?, last_known_json=? WHERE id=?",
                                (now, json.dumps(err_view, ensure_ascii=False), cid))
    finally:
        _SUBMITTING.discard(cid)


async def _refresh_from_gateway(cid: str, gateway_id: str):
    """已经提交给 gateway 的任务，去问一下最新状态，刷新本地记录。"""
    try:
        resp = await client.get(f"{GATEWAY_URL}/v1/videos/{gateway_id}")
    except httpx.HTTPError:
        return
    now = time.time()
    if resp.status_code == 404:
        # gateway 完全不认识这个 gateway_id——通常是所在后端进程被重启（比如手动 flip 分区
        # 或崩溃），内存态清空，这条任务已经不会再有任何进展。不能像别的非 200 一样悄悄跳过，
        # 那样 status 会永远卡在 queued/running，界面看着像在跑，实际已经死了、永远不会再更新。
        with db() as con:
            row = con.execute("SELECT last_known_json FROM jobs WHERE id=?", (cid,)).fetchone()
        if row is None:
            return
        view = json.loads(row["last_known_json"]) if row["last_known_json"] else {"id": cid}
        if view.get("status") in ("done", "failed", "cancelled"):
            return  # 已经是终态，不要覆盖
        view["status"] = "failed"
        view["error"] = "任务在执行途中丢失：所在后端进程被重启，内存态被清空，不会再有进展"
        with db() as con:
            con.execute(
                "UPDATE jobs SET status='failed', updated_at=?, last_known_json=? WHERE id=?",
                (now, json.dumps(view, ensure_ascii=False), cid),
            )
        return
    if resp.status_code != 200:
        return  # 其它非 200（后端一时不可达等）多半是暂时的，保留原状态，下一轮再试
    job = resp.json()
    job["id"] = cid
    with db() as con:
        con.execute(
            "UPDATE jobs SET status=?, backend=?, updated_at=?, last_known_json=? WHERE id=?",
            (job.get("status"), job.get("backend"), now, json.dumps(job, ensure_ascii=False), cid),
        )


async def _discover_from_gateway():
    """兜底：把绕过 console、直接打 GATEWAY_URL（raylight:8100 或老 gateway:8000）提交的任务
    捞回来。两边后端都有 GET /v1/videos 全量列表（raylight 直接返回 list，老 gateway 包一层
    {"jobs":[...]}），本来就是给这个用的。2026-08-24 真出过一次：有人直连后端提交，console
    界面完全看不到，两边还因为都不知道对方在用哪张卡而抢 GPU。
    2026-09-04 更正：本函数确实只在 reconciler_loop 里单协程顺序调用、不会跟自己并发，但会跟
    别的协程（create_video/get_video 触发的 _try_submit）并发——如果正好在某个 _try_submit
    "已经把请求 POST 给 gateway、gateway_id 还没写回数据库"的窗口执行，会把这条其实已经被
    console 自己认领、只是还没记上账的任务当成孤儿，插出一条重复记录（同一个 comfy_prompt_id
    在任务列表里出现两条）。用跟 _try_submit 共享的 _SUBMIT_LOCK 互斥，两段簿记操作不再重叠。"""
    async with _SUBMIT_LOCK:
        try:
            resp = await client.get(f"{GATEWAY_URL}/v1/videos", params={"limit": 200})
        except httpx.HTTPError:
            return
        if resp.status_code != 200:
            return
        body = resp.json()
        remote = body.get("jobs", body) if isinstance(body, dict) else body
        with db() as con:
            known = {r["gateway_id"] for r in con.execute(
                "SELECT gateway_id FROM jobs WHERE gateway_id IS NOT NULL")}
        for summary in remote:
            gid = summary.get("id")
            if not gid or gid in known:
                continue
            try:
                detail = await client.get(f"{GATEWAY_URL}/v1/videos/{gid}")
            except httpx.HTTPError:
                continue
            if detail.status_code != 200:
                continue
            job = detail.json()
            req = job.get("req", {})
            now = time.time()
            job["id"] = gid  # 发现的任务没有独立的 console id，直接用 gateway id 本身
            with db() as con:
                con.execute(
                    "INSERT OR IGNORE INTO jobs "
                    "(id, gateway_id, partition, backend, status, req_json, created_at, updated_at, last_known_json) "
                    "VALUES (?,?,?,?,?,?,?,?,?)",
                    (gid, gid, job.get("partition") or _partition_of(req), job.get("backend"),
                     job.get("status", "queued"), json.dumps(req, ensure_ascii=False),
                     job.get("created_at", now), now, json.dumps(job, ensure_ascii=False)),
                )


@api.post("/v1/videos", status_code=202)
async def create_video(req: dict):
    _check_not_resetting()
    partition = _partition_of(req)
    cid = uuid.uuid4().hex[:12]
    now = time.time()
    synth = {"id": cid, "status": "queued", "req": req, "partition": partition, "created_at": now}
    with db() as con:
        con.execute(
            "INSERT INTO jobs (id, gateway_id, partition, backend, status, req_json, created_at, updated_at, last_known_json) "
            "VALUES (?,NULL,?,NULL,'queued',?,?,?,?)",
            (cid, partition, json.dumps(req, ensure_ascii=False), now, now, json.dumps(synth, ensure_ascii=False)),
        )
    await _try_submit(cid)  # 尽力立即提交一次；提交不上就留给 reconciler_loop
    with db() as con:
        row = con.execute("SELECT * FROM jobs WHERE id=?", (cid,)).fetchone()
    return _row_view(row)


@api.post("/v1/videos/{jid}/cancel")
async def cancel_video(jid: str):
    with db() as con:
        row = con.execute("SELECT * FROM jobs WHERE id=?", (jid,)).fetchone()
    if row is None:
        raise HTTPException(404, "无此任务")
    if row["status"] in ("done", "failed", "cancelled"):
        raise HTTPException(409, f"任务当前状态是 {row['status']}，已经是终态，取消不了")
    now = time.time()
    if row["gateway_id"]:
        was_running = row["status"] == "running"  # 停"running"的任务=硬重启，见下面_mark_resetting
        # 已经提交给后端了，得让后端也知道要取消（后端只允许取消它自己还在排队的任务，
        # 如果已经开始跑了会返回409——这里直接把后端的错误透传出去，不假装取消成功了）
        try:
            resp = await client.post(f"{GATEWAY_URL}/v1/videos/{row['gateway_id']}/cancel")
        except httpx.HTTPError:
            raise HTTPException(502, "取消请求发不到后端，稍后再试")
        if resp.status_code >= 400:
            raise HTTPException(resp.status_code, resp.text)
        if was_running:
            _mark_resetting()
        job = resp.json()
        job["id"] = jid
        with db() as con:
            con.execute(
                "UPDATE jobs SET status='cancelled', updated_at=?, last_known_json=? WHERE id=?",
                (now, json.dumps(job, ensure_ascii=False), jid),
            )
    else:
        # 还没提交给后端（本地排队中），直接在 console 自己这边标掉，reconciler_loop 只挑
        # status='queued' 的行重试提交，改成 cancelled 之后它就不会再去试着提交这条了
        with db() as con:
            con.execute("UPDATE jobs SET status='cancelled', updated_at=? WHERE id=?", (now, jid))
    with db() as con:
        row = con.execute("SELECT * FROM jobs WHERE id=?", (jid,)).fetchone()
    return _row_view(row)


@api.post("/v1/videos/{jid}/prioritize")
async def prioritize_video(jid: str):
    with db() as con:
        row = con.execute("SELECT * FROM jobs WHERE id=?", (jid,)).fetchone()
    if row is None:
        raise HTTPException(404, "无此任务")
    if row["status"] != "queued":
        raise HTTPException(409, f"任务当前状态是 {row['status']}，只有排队中的任务能插队")
    if not row["gateway_id"]:
        raise HTTPException(409, "任务还没提交给后端，稍后再试")
    try:
        resp = await client.post(f"{GATEWAY_URL}/v1/videos/{row['gateway_id']}/prioritize")
    except httpx.HTTPError:
        raise HTTPException(502, "插队请求发不到后端，稍后再试")
    if resp.status_code >= 400:
        raise HTTPException(resp.status_code, resp.text)
    job = resp.json()
    job["id"] = jid
    with db() as con:
        con.execute(
            "UPDATE jobs SET updated_at=?, last_known_json=? WHERE id=?",
            (time.time(), json.dumps(job, ensure_ascii=False), jid),
        )
    return job


@api.get("/v1/videos/{jid}/comfy_history")
async def get_comfy_history(jid: str):
    with db() as con:
        row = con.execute("SELECT gateway_id, backend FROM jobs WHERE id=?", (jid,)).fetchone()
    if row is None:
        raise HTTPException(404, "无此任务")
    if not row["gateway_id"] or not row["backend"]:
        raise HTTPException(409, "任务还在本地排队，尚未分配到可用后端")
    try:
        resp = await client.get(f"{row['backend']}/v1/videos/{row['gateway_id']}/comfy_history")
    except httpx.HTTPError:
        raise HTTPException(502, "查询 ComfyUI history 发不到后端，稍后再试")
    if resp.status_code >= 400:
        raise HTTPException(resp.status_code, resp.text)
    return resp.json()


@api.get("/v1/videos/{jid}")
async def get_video(jid: str):
    with db() as con:
        row = con.execute("SELECT * FROM jobs WHERE id=?", (jid,)).fetchone()
    if row is None:
        raise HTTPException(404, "无此任务")
    if row["gateway_id"]:
        await _refresh_from_gateway(jid, row["gateway_id"])
    else:
        await _try_submit(jid)  # 还在本地排队时，顺手趁轮询这一下也试着提交
    with db() as con:
        row = con.execute("SELECT * FROM jobs WHERE id=?", (jid,)).fetchone()
    return _row_view(row)


@api.get("/v1/videos/{jid}/file")
async def download_video(jid: str, request: Request):
    with db() as con:
        row = con.execute(
            "SELECT gateway_id, backend, status, last_known_json FROM jobs WHERE id=?", (jid,)
        ).fetchone()
    if row is None:
        raise HTTPException(404, "无此任务")
    if not row["gateway_id"] or not row["backend"]:
        raise HTTPException(409, "任务还在本地排队，尚未分配到可用后端")

    # 不经过 gateway.py 转发：它自己那层 `StreamingResponse(urlopen(...))` 不认 Range 请求，
    # 探测 Range 一律回 200 而不是 206，还不带 Content-Length/Accept-Ranges——`<video>`
    # 标签播放前先探一次 Range 拿时长/可否拖动，探不到就一直卡在 readyState=0 播不出来
    # （下载能成功是因为普通 GET 不需要这些，只有网页里的播放器会被这个坑绊住）。
    # api.py 自己那层是 Starlette `FileResponse`，原生支持 Range；console 既然已经从
    # gateway 的响应里存了 `backend`（真实拥有这条任务的那个 GPU 后端），直接绕过 gateway
    # 这层转发、把 Range 头透传给 api.py，把它的状态码/头原样带回去，问题在这一层就解决，
    # 不用碰 gateway.py/api.py 一行代码。
    # 下载/播放到本地的文件名——不带这个的话浏览器只能从 URL 最后一段猜，我们的 URL 都是
    # 以 "/file" 结尾（没有扩展名），猜出来清一色叫 "file.mp4"，多个任务的视频存到同一个
    # 下载目录会互相覆盖。用 console 自己的任务 id 拼一个每条任务都不同的文件名。
    download_filename = f"h3-{jid}.mp4"
    content_disposition = f'attachment; filename="{download_filename}"'

    upstream_url = f"{row['backend']}/v1/videos/{row['gateway_id']}/file"
    headers = {}
    if "range" in request.headers:
        headers["range"] = request.headers["range"]
    try:
        req = client.build_request("GET", upstream_url, headers=headers)
        upstream = await client.send(req, stream=True)
    except httpx.HTTPError:
        upstream = None

    if upstream is None or upstream.status_code >= 400:
        if upstream is not None:
            await upstream.aclose()
        # 后端进程重启会把 api.py 内存里的 JOBS 字典清空（这是它本来的设计，任务状态
        # 不落盘），导致它忘了这条任务——但视频文件本身在磁盘上没被动过。console 和
        # api.py 跑在同一台机器上，任务完成时已经把 `file` 路径存进了自己的 sqlite，
        # 后端忘了就直接读文件兜底，不用等它想起来（它也不会想起来，重启就是清空）。
        cached = json.loads(row["last_known_json"]) if row["last_known_json"] else {}
        file_path = cached.get("file")
        if file_path and os.path.exists(file_path):
            return FileResponse(file_path, media_type="video/mp4",
                                 headers={"Content-Disposition": content_disposition})
        raise HTTPException(404, "后端进程重启后忘了这条任务的记录，磁盘上也没找到对应文件")

    async def stream():
        async for chunk in upstream.aiter_bytes():
            yield chunk
        await upstream.aclose()

    passthrough = {k: v for k, v in upstream.headers.items()
                   if k.lower() in ("content-length", "content-range", "accept-ranges", "content-type")}
    passthrough["content-disposition"] = content_disposition
    return StreamingResponse(stream(), status_code=upstream.status_code, headers=passthrough)


@api.get("/v1/videos")
async def list_videos(limit: int = 20, offset: int = 0, trash: bool = False, q: str = "", status: str = "",
                       kind: str = "", sort: str = "desc"):
    # console 自己的 sqlite 是任务列表的唯一真相来源（本地排队中的任务，gateway 根本不知道它存在，
    # 没法从 gateway 反查出来）。实时进度由后台 reconciler_loop 持续刷新进 last_known_json，
    # 这里直接读，不在请求路径里再打一遍 gateway，避免列表接口被拖慢。
    # 软删除：deleted_at 非空的任务默认不出现在列表里，?trash=true 反过来只看这些——
    # 数据还在同一张表里，删除/恢复就是 UPDATE 一个时间戳，不用挪表、不用真删文件。
    where = ["deleted_at IS NOT NULL" if trash else "deleted_at IS NULL"]
    params: list = []
    if q:
        # prompt 只存在 req_json 这坨 JSON 文本里，没有单独的列——json.dumps 是
        # ensure_ascii=False 存的，中文原样落盘，直接 LIKE 子串匹配就够用，不用上 json_extract。
        where.append("(id LIKE ? OR req_json LIKE ?)")
        params += [f"%{q}%", f"%{q}%"]
    if status:
        where.append("status = ?")
        params.append(status)
    if kind:
        # 前端叫它"类型"（文生/首尾帧/参考图），数据库列叫 partition——是同一个值
        # （见 _partition_of），换个外部参数名单纯是避免跟 /v1/backends 里 GPU 分区的
        # "partition" 混着看。
        where.append("partition = ?")
        params.append(kind)
    order_col = "deleted_at" if trash else "created_at"
    order = f"{order_col} {'ASC' if sort == 'asc' else 'DESC'}"
    with db() as con:
        # 前端要做真分页（第X/共Y页），得知道筛选条件下总共有多少条，不能只看这一页拿回了几条。
        total = con.execute(f"SELECT COUNT(*) FROM jobs WHERE {' AND '.join(where)}", params).fetchone()[0]
        rows = con.execute(
            f"SELECT * FROM jobs WHERE {' AND '.join(where)} ORDER BY {order} LIMIT ? OFFSET ?",
            params + [limit, offset],
        ).fetchall()
    return {"jobs": [_row_view(r) for r in rows], "total": total}


@api.delete("/v1/videos/{jid}")
async def delete_video(jid: str):
    """软删除：只挪出主列表，不碰 gateway/api.py 上的任务本身（可能还在跑），不删磁盘文件。"""
    with db() as con:
        cur = con.execute("UPDATE jobs SET deleted_at=? WHERE id=? AND deleted_at IS NULL", (time.time(), jid))
    if cur.rowcount == 0:
        raise HTTPException(404, "无此任务，或已经在回收站里")
    return {"id": jid, "deleted": True}


@api.post("/v1/videos/{jid}/restore")
async def restore_video(jid: str):
    with db() as con:
        cur = con.execute("UPDATE jobs SET deleted_at=NULL WHERE id=? AND deleted_at IS NOT NULL", (jid,))
    if cur.rowcount == 0:
        raise HTTPException(404, "无此任务，或不在回收站里")
    return {"id": jid, "deleted": False}


@api.get("/v1/backends")
async def backends():
    resp = await client.get(f"{GATEWAY_URL}/health")
    resp.raise_for_status()
    return resp.json()


@api.post("/v1/admin/reset-gpu")
async def reset_gpu():
    # 通用重置入口，不挂在具体任务上——应对Raylight actor崩溃/OOM后不自愈、显存卡在接近
    # 满载但JOBS里没有running任务的情况（这时候"停止"按钮不会出现，需要这个兜底）。
    try:
        resp = await client.post(f"{GATEWAY_URL}/admin/reset-gpu")
    except httpx.HTTPError:
        raise HTTPException(502, "重置请求发不到后端，稍后再试")
    if resp.status_code >= 400:
        raise HTTPException(resp.status_code, resp.text)
    _mark_resetting()
    return resp.json()


@api.get("/v1/models")
async def models():
    # partitions 是静态信息，对应 raylight/server.py 里的 DIT_FL2VA/DIT_REF2VA 常量，只读展示。
    # 2026-09-03：presets 之前是硬编码的3档旧数据（模型文件名对应的是已删除的 api.py/gateway.py
    # 那套旧架构，跟 raylight 现在实际用的模型对不上），改成实时问 raylight 的
    # /v1/resolutions（RESOLUTION_CEILINGS 的镜像），保证跟实际支持的分辨率档位一致，
    # 不用两边手动同步一份数字。
    try:
        resp = await client.get(f"{GATEWAY_URL}/v1/resolutions")
        resolutions = resp.json()["resolutions"] if resp.status_code == 200 else []
    except httpx.HTTPError:
        resolutions = []
    presets = [
        {"label": f"{r['width']}×{r['height']}，最多 {r['max_frames']} 帧（≈{r['max_frames']/24:.1f}s@24fps）",
         "width": r["width"], "height": r["height"], "max_frames": r["max_frames"]}
        for r in resolutions
    ]
    return {
        "partitions": {
            "fl2va": {"dit": "minimax_h3_fl2va_pruned_int8_convrot.safetensors",
                      "desc": "首末帧条件生成"},
            "ref2va": {"dit": "minimax_h3_ref2va_pruned_int8_convrot.safetensors",
                       "desc": "文生视频 / 多参考图生成"},
        },
        "presets": presets,
    }


# ---------------- 弹性 partition 调度器（仅服务器部署时启用） ----------------
def _pending_counts() -> dict:
    with db() as con:
        rows = con.execute(
            "SELECT partition, COUNT(*) c FROM jobs WHERE status IN ('queued','running') GROUP BY partition"
        ).fetchall()
    return {r["partition"]: r["c"] for r in rows}


async def _flip_backend(gpu: int, port: int, new_partition: str):
    # 和今天手动测试时用的手法一致：按端口查 PID、杀、用现成的 api-run.sh 重启到新分区。
    # 不新写进程管理代码，复用已经在跑的脚本。
    # 注意：ss 的输出里 pid 不是独立的空格分隔 token，是嵌在 `users:(("python",pid=NNNN,fd=14))`
    # 里面的——按空格 split 后拿 `tok.startswith("pid=")` 永远不会命中，得用正则从整段里抠出来
    # （手动测试时用的是 `grep -o "pid=[0-9]*"`，这里必须做等价的事，不能简化成 split）。
    pid_out = subprocess.run(["ss", "-lptn", f"sport = :{port}"], capture_output=True, text=True).stdout
    for pid in re.findall(r"pid=(\d+)", pid_out):
        subprocess.run(["kill", pid])
    subprocess.Popen(
        ["setsid", "nohup", "./api-run.sh", str(gpu), str(port), new_partition],
        cwd=str(H3_ROOT), stdout=open(H3_ROOT / "logs" / f"api-{port}.log", "a"),
        stderr=subprocess.STDOUT, stdin=subprocess.DEVNULL, start_new_session=True,
    )


async def supervisor_loop():
    idle_since: dict[int, float] = {}
    # 一个 flip 从 kill 到新进程报 ready 要 13-20 秒（api.py 冷启动实测），远长于 10 秒的 tick 间隔。
    # 如果不记「已经在切换中」，接下来两三轮 tick 会看到 fl2va_ready 还是 0（新进程没起来），
    # 误判成需求仍未满足，再多切几张卡——今天调试时真实撞到过：一个 pending=1 的请求把三张卡都切了
    # 过去（而且当时 kill 那步还有个 bug 没真的杀掉旧进程，新进程绑端口失败悄悄退出，见下）。
    flipping: dict[int, tuple[str, float]] = {}  # port -> (目标 partition, 发起时间)
    while True:
        await asyncio.sleep(10)
        try:
            health = (await client.get(f"{GATEWAY_URL}/health")).json()["backends"]
            pending = _pending_counts()
            now = time.time()

            for port in list(flipping):
                target, started = flipping[port]
                b = next((x for x in health if x["backend"].endswith(f":{port}")), None)
                if b and b.get("partition") == target and b.get("status") == "ok":
                    flipping.pop(port)  # 切换完成，新分区已经在 health 里报 ready 了
                elif now - started > 60:
                    flipping.pop(port)  # 60 秒还没起来，当作这次尝试失败，让下一轮按实际情况重新决定

            fl2va_ready = [b for b in health if b.get("partition") == "fl2va" and b.get("status") == "ok"]
            ref2va_ready = [b for b in health if b.get("partition") == "ref2va" and b.get("status") == "ok"]
            fl2va_inflight = sum(1 for t, _ in flipping.values() if t == "fl2va")
            ref2va_inflight = sum(1 for t, _ in flipping.values() if t == "ref2va")

            def idle_backends(group):
                out = []
                for b in group:
                    port = int(b["backend"].rsplit(":", 1)[1])
                    if port in flipping:
                        continue  # 已经在切换中，不能再选它
                    if b.get("queued", 0) == 0 and b.get("running", 0) == 0:
                        idle_since.setdefault(port, now)
                        if now - idle_since[port] >= 30:
                            out.append(port)
                    else:
                        idle_since.pop(port, None)
                return out

            fl2va_pending = pending.get("fl2va", 0)
            ref2va_pending = pending.get("ref2va", 0)
            print(f"[supervisor] tick: pending={pending} "
                  f"fl2va={len(fl2va_ready)}(+{fl2va_inflight} 切换中) "
                  f"ref2va={len(ref2va_ready)}(+{ref2va_inflight} 切换中)", flush=True)

            if fl2va_pending > len(fl2va_ready) + fl2va_inflight:
                candidates = idle_backends(ref2va_ready)
                if candidates:
                    port = candidates[0]
                    gpu = next(t["gpu"] for t in BACKEND_TOPOLOGY if t["port"] == port)
                    print(f"[supervisor] flip port={port} gpu={gpu} ref2va -> fl2va", flush=True)
                    flipping[port] = ("fl2va", now)
                    idle_since.pop(port, None)
                    await _flip_backend(gpu, port, "fl2va")
            elif ref2va_pending > len(ref2va_ready) + ref2va_inflight:
                candidates = idle_backends(fl2va_ready)
                if candidates:
                    port = candidates[0]
                    gpu = next(t["gpu"] for t in BACKEND_TOPOLOGY if t["port"] == port)
                    print(f"[supervisor] flip port={port} gpu={gpu} fl2va -> ref2va", flush=True)
                    flipping[port] = ("ref2va", now)
                    idle_since.pop(port, None)
                    await _flip_backend(gpu, port, "ref2va")
        except Exception as e:
            print(f"[supervisor] tick failed: {e}", flush=True)


# ---------------- 常驻对账循环（本地/服务器都跑，跟 GPU 无关） ----------------
# 两件事：本地排队中的任务（gateway_id 还是空）隔几秒重试提交一次；已提交的任务刷新最新状态。
# 这是任务列表能实时更新、以及「没后端时先排队而不是直接报错」这两件事真正生效的地方。
async def reconciler_loop():
    while True:
        await asyncio.sleep(5)
        try:
            with db() as con:
                pending = con.execute(
                    "SELECT id FROM jobs WHERE gateway_id IS NULL AND status='queued'"
                ).fetchall()
                in_flight = con.execute(
                    "SELECT id, gateway_id FROM jobs WHERE gateway_id IS NOT NULL AND status IN ('queued','running')"
                ).fetchall()
            for r in pending:
                await _try_submit(r["id"])
            for r in in_flight:
                await _refresh_from_gateway(r["id"], r["gateway_id"])
            await _discover_from_gateway()
        except Exception as e:
            print(f"[reconciler] tick failed: {e}", flush=True)


@app.on_event("startup")
async def _startup():
    asyncio.create_task(reconciler_loop())
    if ENABLE_SUPERVISOR:
        asyncio.create_task(supervisor_loop())
        print("[console] elastic partition supervisor enabled", flush=True)
    else:
        print("[console] supervisor disabled (H3_ENABLE_SUPERVISOR!=1) — fl2va/ref2va allocation is whatever the backends currently report", flush=True)


# 任务详情/新建各自有一个真实的 URL，方便复制粘贴分享、刷新不丢状态——都是同一个单页应用
# （static/index.html 自己那份 JS 读 location.pathname 决定显示哪个视图），这三条路由只是让
# 浏览器地址栏里出现符合预期的路径，不是三个不同的页面文件。
#
# 页面不再直接 FileResponse 原文件，而是读进来做一次字符串替换：把 __H3_BASE_PATH__ 占位符
# 换成 BASE_PATH。部署在共享域名子路径下时（比如 qwen.wudaima.com/console），前端 JS 里所有
# fetch/history.pushState 用的绝对路径都得带上这个前缀，否则请求会打到域名根路径而不是我们
# 这个子路径——BASE_PATH 留空时替换成空字符串，跟原来的行为完全一样。
def render_index() -> HTMLResponse:
    html = (BASE / "static" / "index.html").read_text(encoding="utf-8")
    return HTMLResponse(html.replace("__H3_BASE_PATH__", BASE_PATH))


@router.get("/")
async def index_page():
    return render_index()


@router.get("/task.html/{tid}")
async def task_page(tid: str):
    return render_index()


@router.get("/newtask.html")
async def newtask_page():
    return render_index()


# BASE_PATH 留空时这个前缀是空字符串，router 上的路径原样挂到 app 根上，行为跟以前
# 直接用 @app.xxx 一模一样；非空时（比如 "/console"）整体套一层前缀，一次性生效，
# 不用每个路由自己拼。
app.include_router(router, prefix=BASE_PATH)
app.include_router(api, prefix=BASE_PATH)