modelbeast/server/main.py

import asyncio
import json
from pathlib import Path

from fastapi import FastAPI, File, HTTPException, UploadFile, WebSocket, WebSocketDisconnect
from fastapi.responses import FileResponse, JSONResponse
from fastapi.staticfiles import StaticFiles

from . import db, registry, settings as settings_mod, store
from .runner import runner

app = FastAPI(title="MODELBEAST")
WEB_DIST = db.ROOT / "web" / "dist"
INBOX = db.DATA / "inbox"


@app.on_event("startup")
async def startup():
    runner.operators = registry.load_operators()
    con = db.connect()
    app.state.con = con
    runner.get_settings = lambda: settings_mod.get_all(con)
    asyncio.create_task(runner.worker())
    asyncio.create_task(watch_inbox())


# -- operators ----------------------------------------------------------------
@app.get("/api/operators")
def list_operators():
    return [{k: v for k, v in op.items() if k != "dir"} for op in runner.operators.values()]


# -- settings -----------------------------------------------------------------
@app.get("/api/settings")
def get_settings():
    s = settings_mod.get_all(app.state.con)
    view = settings_mod.public_view(s)
    view["_secret_keys"] = sorted(settings_mod.SECRET_KEYS)
    view["_env_keys"] = sorted(settings_mod.ENV_MAP.keys())
    # env var names currently satisfied — lets the UI gate operators
    view["_env_set"] = sorted(env for key, env in settings_mod.ENV_MAP.items() if s.get(key))
    return view


@app.put("/api/settings")
def put_settings(payload: dict):
    # ignore masked secret placeholders so re-saving the form doesn't wipe a key
    clean = {k: v for k, v in payload.items()
             if not (k in settings_mod.SECRET_KEYS and set(str(v)) <= {"•"})}
    settings_mod.set_many(app.state.con, clean)
    return get_settings()


# -- assets ---------------------------------------------------------------------
@app.get("/api/assets")
def list_assets():
    return store.list_assets(app.state.con)


@app.post("/api/assets")
async def upload_asset(file: UploadFile = File(...)):
    data = await file.read()
    if not data:
        raise HTTPException(400, "empty file")
    return store.register_upload(app.state.con, file.filename or "unnamed", data)


@app.delete("/api/assets/{asset_id}")
def delete_asset(asset_id: str):
    if not store.delete_asset(app.state.con, asset_id):
        raise HTTPException(404, "no such asset")
    return {"ok": True}


@app.get("/api/assets/{asset_id}/file")
def asset_file(asset_id: str, member: str | None = None):
    asset = store.get_asset(app.state.con, asset_id)
    if not asset:
        raise HTTPException(404, "no such asset")
    path = Path(asset["path"])
    if member:
        target = (path / member).resolve()
        if not str(target).startswith(str(path.resolve())) or not target.is_file():
            raise HTTPException(404, "no such member")
        return FileResponse(target)
    if path.is_dir():
        files = sorted(p.name for p in path.iterdir() if p.is_file())
        return JSONResponse({"folder": True, "files": files})
    return FileResponse(path, filename=asset["name"])


# -- jobs -----------------------------------------------------------------------
@app.get("/api/jobs")
def list_jobs():
    return runner.list_jobs(app.state.con)


@app.post("/api/jobs")
async def create_job(payload: dict):
    operator = payload.get("operator")
    if operator not in runner.operators:
        raise HTTPException(400, f"unknown operator: {operator}")
    asset_ids = payload.get("asset_ids")
    if not asset_ids:
        asset_ids = [payload["asset_id"]] if payload.get("asset_id") else []
    for aid in asset_ids:
        if not store.get_asset(app.state.con, aid):
            raise HTTPException(400, f"unknown asset: {aid}")
    params = payload.get("params") or {}
    job = runner.create_job(app.state.con, operator, asset_ids, params)
    await runner.broadcast({"type": "job", "job": job})
    return job


@app.get("/api/jobs/{job_id}")
def get_job(job_id: str):
    job = runner.get_job(app.state.con, job_id)
    if not job:
        raise HTTPException(404, "no such job")
    return job


@app.post("/api/jobs/{job_id}/cancel")
async def cancel_job(job_id: str):
    if not await runner.cancel(app.state.con, job_id):
        raise HTTPException(400, "job not cancellable")
    return {"ok": True}


@app.post("/api/jobs/{job_id}/retry")
async def retry_job(job_id: str):
    job = runner.get_job(app.state.con, job_id)
    if not job:
        raise HTTPException(404, "no such job")
    new = runner.create_job(app.state.con, job["operator"],
                            job.get("asset_ids") or [], job["params"])
    await runner.broadcast({"type": "job", "job": new})
    return new


@app.delete("/api/jobs/{job_id}")
def delete_job(job_id: str):
    if not runner.delete_job(app.state.con, job_id):
        raise HTTPException(400, "job running or missing")
    return {"ok": True}


# -- inbox watch folder ---------------------------------------------------------
async def watch_inbox():
    """Ingest files dropped into data/inbox/ once their size is stable (avoids
    grabbing partial copies). Lets UE/Blender/rsync drop outputs in for pickup."""
    INBOX.mkdir(parents=True, exist_ok=True)
    seen: dict[str, int] = {}
    while True:
        await asyncio.sleep(3)
        try:
            for p in INBOX.iterdir():
                if p.name.startswith(".") or not p.exists():
                    continue
                size = p.stat().st_size if p.is_file() else sum(
                    f.stat().st_size for f in p.rglob("*") if f.is_file())
                if seen.get(str(p)) == size and size > 0:
                    asset = store.register_file(app.state.con, p, move=True,
                                                meta={"source": "inbox"})
                    seen.pop(str(p), None)
                    await runner.broadcast({"type": "assets_changed", "assets": [asset["id"]]})
                else:
                    seen[str(p)] = size
        except Exception as e:
            print(f"[inbox] {e}")


# -- websocket ---------------------------------------------------------------
@app.websocket("/ws")
async def ws(websocket: WebSocket):
    await websocket.accept()
    runner.subscribers.add(websocket)
    try:
        while True:
            await websocket.receive_text()
    except WebSocketDisconnect:
        runner.subscribers.discard(websocket)


# -- frontend (built) -----------------------------------------------------------
if WEB_DIST.exists():
    app.mount("/", StaticFiles(directory=WEB_DIST, html=True), name="web")