ubuntu-bare-metal/stacks/server-ui/app.py

#!/usr/bin/env python3
"""GMKtec K11 — Docker stack management UI."""

from __future__ import annotations

import logging
import os
import sys
from pathlib import Path
from typing import Any

import uvicorn
from fastapi import Depends, FastAPI, HTTPException, Query, Request, WebSocket
from fastapi.responses import FileResponse, Response
from fastapi.staticfiles import StaticFiles

from pydantic import BaseModel, Field

from compose_runner import (
    ComposeError,
    PortConfigError,
    StackConfig,
    find_compose_file,
    find_running_gpu_stacks,
    get_container_state,
    get_logs,
    load_stacks,
    set_stack_port,
    stack_restart,
    stack_start,
    stack_stop,
    stack_update,
)
from cli_pty import (
    CliSessionLimitError,
    get_session_manager,
    init_session_manager,
    run_pty_session,
    sessions_available,
)
from comfyui_models import (
    ComfyUIModelError,
    DeleteModelError,
    JobConflictError,
    JobNotFoundError,
    ModelNotFoundError,
    disk_usage,
    get_download_manager,
    init_download_manager,
    list_installed,
)
from file_explorer import FileExplorer, FileExplorerError
from gpu_info import GpuInfoError, query_gpu
from gpu_fan_proxy import AgentProxyError, check_agent_health, forward_request, resolve_agent_url

STACK_DIR = Path(__file__).resolve().parent

# Unified control-plane env (see stacks/control-plane/env_loader.py)
_local_env_loader = STACK_DIR / "env_loader.py"
if _local_env_loader.is_file():
    if str(STACK_DIR) not in sys.path:
        sys.path.insert(0, str(STACK_DIR))
else:
    for _cp in (
        "/opt/control-plane",
        "/repo/stacks/control-plane",
        str(STACK_DIR.parent / "control-plane"),
    ):
        if _cp not in sys.path and Path(_cp).joinpath("env_loader.py").is_file():
            sys.path.insert(0, _cp)
from env_loader import api_key_source, load_control_plane_env, resolve_repo_root  # noqa: E402
STATIC_DIR = STACK_DIR / "static"
CONFIG_PATH = STACK_DIR / "stacks.yaml"

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(levelname)s %(message)s",
    handlers=[logging.StreamHandler(sys.stdout)],
)
log = logging.getLogger(__name__)


ENV = load_control_plane_env(STACK_DIR)
HOST = ENV.get("SERVER_UI_HOST", "0.0.0.0")
PORT = int(ENV.get("SERVER_UI_PORT", "8091"))
API_KEY = ENV.get("API_KEY", "")
GPU_FAN_AGENT_URL = resolve_agent_url(
    ENV.get("GPU_FAN_AGENT_URL", "http://127.0.0.1:18090"),
    API_KEY,
)
REPO_ROOT = resolve_repo_root(STACK_DIR, ENV)
STACKS_BASE = REPO_ROOT / "stacks"

STACKS: list[StackConfig] = load_stacks(CONFIG_PATH, STACKS_BASE)
STACK_BY_ID: dict[str, StackConfig] = {s.id: s for s in STACKS}

FILE_EXPLORER_ROOT = ENV.get("FILE_EXPLORER_ROOT", "/")
FILE_EXPLORER_MAX_BYTES = int(ENV.get("FILE_EXPLORER_MAX_BYTES", "2097152"))
file_explorer = FileExplorer(root=FILE_EXPLORER_ROOT, max_bytes=FILE_EXPLORER_MAX_BYTES)

CLI_ENABLED = ENV.get("CLI_ENABLED", "1").strip().lower() in ("1", "true", "yes")
CLI_SHELL = ENV.get("CLI_SHELL", "/bin/bash")
CLI_DEFAULT_CWD = ENV.get("CLI_DEFAULT_CWD", os.path.expanduser("~"))
CLI_MAX_SESSIONS = int(ENV.get("CLI_MAX_SESSIONS", "5"))
CLI_SESSION_TTL = int(ENV.get("CLI_SESSION_TTL", "300"))

DATA_ROOT = Path(ENV.get("DATA_ROOT", "/data"))
COMFYUI_JOB_TTL = int(ENV.get("COMFYUI_JOB_TTL", "3600"))

init_session_manager(session_ttl=CLI_SESSION_TTL, max_sessions=CLI_MAX_SESSIONS)
init_download_manager(REPO_ROOT, DATA_ROOT, job_ttl=COMFYUI_JOB_TTL)

app = FastAPI(title="Server UI", version="1.0.0")
app.mount("/static", StaticFiles(directory=STATIC_DIR), name="static")


class PortUpdate(BaseModel):
    port: int = Field(ge=1024, le=65535)
    recreate: bool = True


class FileWriteBody(BaseModel):
    path: str
    content: str = ""


class FilePathBody(BaseModel):
    path: str


class FileRenameBody(BaseModel):
    old_path: str
    new_path: str


class AuthVerifyBody(BaseModel):
    api_key: str = ""


class ComfyUIDownloadBody(BaseModel):
    model_id: str = Field(min_length=1)


class ComfyUIFileDownloadBody(BaseModel):
    url: str = Field(min_length=1)
    relative_path: str = Field(min_length=1)


def get_stack(stack_id: str) -> StackConfig:
    stack = STACK_BY_ID.get(stack_id)
    if not stack:
        raise HTTPException(status_code=404, detail=f"Unknown stack: {stack_id}")
    return stack


def require_mutation_auth(request: Request) -> None:
    if not API_KEY:
        return
    key = request.headers.get("X-API-Key", "")
    if key != API_KEY:
        raise HTTPException(status_code=401, detail="Invalid or missing API key")


def _is_lan_bind() -> bool:
    return HOST.strip().lower() not in ("127.0.0.1", "localhost", "::1")


def require_gpu_fan_read_auth(request: Request) -> None:
    if not _is_lan_bind() or not API_KEY:
        return
    require_mutation_auth(request)


def require_file_auth(request: Request) -> None:
    if not _is_lan_bind() or not API_KEY:
        return
    require_mutation_auth(request)


def require_comfyui_auth(request: Request) -> None:
    require_file_auth(request)


def _comfyui_exc_to_http(exc: ComfyUIModelError) -> HTTPException:
    if isinstance(exc, ModelNotFoundError):
        return HTTPException(status_code=404, detail=str(exc))
    if isinstance(exc, JobNotFoundError):
        return HTTPException(status_code=404, detail=str(exc))
    if isinstance(exc, JobConflictError):
        return HTTPException(status_code=409, detail=str(exc))
    if isinstance(exc, DeleteModelError):
        return HTTPException(status_code=403, detail=str(exc))
    return HTTPException(status_code=500, detail=str(exc))


def _cli_ws_auth_ok(websocket: WebSocket) -> bool:
    if not _is_lan_bind() or not API_KEY:
        return True
    return websocket.query_params.get("api_key", "") == API_KEY


def _file_exc_to_http(exc: FileExplorerError) -> HTTPException:
    msg = str(exc)
    if "za duż" in msg.lower() or "za duża" in msg.lower():
        return HTTPException(status_code=413, detail=msg)
    if "poza dozwolonym" in msg.lower():
        return HTTPException(status_code=403, detail=msg)
    if "Brak uprawnień" in msg:
        return HTTPException(status_code=403, detail=msg)
    if "nie jest pusty" in msg.lower() or "już istnieje" in msg.lower():
        return HTTPException(status_code=409, detail=msg)
    if "Nie znaleziono" in msg:
        return HTTPException(status_code=404, detail=msg)
    if "nie jest katalogiem" in msg.lower() or "To jest katalog" in msg:
        return HTTPException(status_code=400, detail=msg)
    return HTTPException(status_code=400, detail=msg)


def check_gpu_policy(stack: StackConfig) -> None:
    if not stack.gpu:
        return
    conflicts = find_running_gpu_stacks(STACKS, exclude_id=stack.id)
    if conflicts:
        names = ", ".join(conflicts)
        raise HTTPException(
            status_code=409,
            detail=f"GPU conflict: stop running GPU stack(s) first: {names}",
        )


@app.get("/")
def index() -> FileResponse:
    return FileResponse(STATIC_DIR / "index.html")


@app.get("/api/health")
def api_health() -> dict[str, bool]:
    return {"ok": True}


@app.post("/api/auth/verify")
def api_auth_verify(body: AuthVerifyBody) -> dict[str, bool]:
    if not API_KEY:
        return {"ok": True}
    if body.api_key == API_KEY:
        return {"ok": True}
    raise HTTPException(status_code=401, detail="Invalid or missing API key")


@app.get("/api/gpu")
def api_gpu() -> dict[str, Any]:
    try:
        return query_gpu()
    except GpuInfoError as exc:
        raise HTTPException(status_code=503, detail=str(exc)) from exc


@app.get("/api/stacks")
def api_stacks() -> dict[str, Any]:
    items = [get_container_state(s) for s in STACKS]
    gpu_running = [i["id"] for i in items if i["gpu"] and i["running"]]
    return {
        "stacks": items,
        "gpu_running": gpu_running,
        "repo_root": str(REPO_ROOT),
    }


@app.get("/api/stacks/{stack_id}/logs")
def api_logs(stack_id: str, tail: int = Query(default=100, ge=1, le=500)) -> dict[str, str]:
    stack = get_stack(stack_id)
    try:
        text = get_logs(stack, tail=tail)
    except ComposeError as exc:
        raise HTTPException(status_code=500, detail=str(exc)) from exc
    return {"stack_id": stack_id, "logs": text}


@app.post("/api/stacks/{stack_id}/start")
def api_start(stack_id: str, _: None = Depends(require_mutation_auth)) -> dict[str, Any]:
    stack = get_stack(stack_id)
    check_gpu_policy(stack)
    try:
        output = stack_start(stack)
    except ComposeError as exc:
        raise HTTPException(status_code=500, detail=str(exc)) from exc
    return {"ok": True, "action": "start", "stack_id": stack_id, "output": output}


@app.post("/api/stacks/{stack_id}/stop")
def api_stop(stack_id: str, _: None = Depends(require_mutation_auth)) -> dict[str, Any]:
    stack = get_stack(stack_id)
    try:
        output = stack_stop(stack)
    except ComposeError as exc:
        raise HTTPException(status_code=500, detail=str(exc)) from exc
    return {"ok": True, "action": "stop", "stack_id": stack_id, "output": output}


@app.post("/api/stacks/{stack_id}/restart")
def api_restart(stack_id: str, _: None = Depends(require_mutation_auth)) -> dict[str, Any]:
    stack = get_stack(stack_id)
    if stack.gpu:
        conflicts = find_running_gpu_stacks(STACKS, exclude_id=stack.id)
        if conflicts:
            check_gpu_policy(stack)
    try:
        output = stack_restart(stack)
    except ComposeError as exc:
        raise HTTPException(status_code=500, detail=str(exc)) from exc
    return {"ok": True, "action": "restart", "stack_id": stack_id, "output": output}


@app.post("/api/stacks/{stack_id}/update")
def api_update(stack_id: str, _: None = Depends(require_mutation_auth)) -> dict[str, Any]:
    stack = get_stack(stack_id)
    try:
        output = stack_update(stack)
    except ComposeError as exc:
        raise HTTPException(status_code=500, detail=str(exc)) from exc
    return {"ok": True, "action": "update", "stack_id": stack_id, "output": output}


@app.patch("/api/stacks/{stack_id}/port")
def api_set_port(
    stack_id: str,
    body: PortUpdate,
    _: None = Depends(require_mutation_auth),
) -> dict[str, Any]:
    stack = get_stack(stack_id)
    if not stack.port_editable or not stack.port_env:
        raise HTTPException(
            status_code=400,
            detail=f"Port is not editable for stack: {stack_id}",
        )
    try:
        return set_stack_port(
            stack,
            body.port,
            STACKS,
            PORT,
            recreate=body.recreate,
        )
    except PortConfigError as exc:
        raise HTTPException(status_code=409, detail=str(exc)) from exc
    except ComposeError as exc:
        raise HTTPException(status_code=500, detail=str(exc)) from exc


@app.get("/api/gpu-fan/health")
def api_gpu_fan_health(_: None = Depends(require_gpu_fan_read_auth)) -> dict[str, Any]:
    return check_agent_health(GPU_FAN_AGENT_URL, API_KEY)


@app.api_route("/api/gpu-fan/{path:path}", methods=["GET", "PUT", "POST"])
async def api_gpu_fan_proxy(
    path: str,
    request: Request,
    _: None = Depends(require_gpu_fan_read_auth),
) -> Response:
    if request.method in ("PUT", "POST"):
        require_mutation_auth(request)

    body = await request.body()
    try:
        status, resp_headers, payload = forward_request(
            GPU_FAN_AGENT_URL,
            request.method,
            path,
            body=body if body else None,
            agent_api_key=API_KEY,
            content_type=request.headers.get("content-type", "application/json"),
        )
    except AgentProxyError as exc:
        raise HTTPException(
            status_code=502,
            detail=f"GPU fan agent unavailable at {GPU_FAN_AGENT_URL}: {exc}",
        ) from exc

    media_type = resp_headers.get("Content-Type", "application/json")
    return Response(content=payload, status_code=status, media_type=media_type)


@app.get("/api/comfyui/models")
def api_comfyui_models(_: None = Depends(require_comfyui_auth)) -> dict[str, Any]:
    mgr = get_download_manager()
    try:
        payload = mgr.list_models_response()
        payload["disk"] = disk_usage(mgr.data_root)
    except ComfyUIModelError as exc:
        raise _comfyui_exc_to_http(exc) from exc
    return payload


@app.get("/api/comfyui/models/installed")
def api_comfyui_models_installed(_: None = Depends(require_comfyui_auth)) -> dict[str, Any]:
    mgr = get_download_manager()
    return {"installed": list_installed(mgr.data_root), "data_root": str(mgr.data_root)}


@app.post("/api/comfyui/models/download")
def api_comfyui_models_download(
    body: ComfyUIDownloadBody,
    _: None = Depends(require_mutation_auth),
) -> dict[str, Any]:
    mgr = get_download_manager()
    model_id = body.model_id.strip()
    try:
        if model_id.startswith("wf:"):
            job = mgr.start_workflow_download(model_id)
        else:
            job = mgr.start_download(model_id)
    except ComfyUIModelError as exc:
        raise _comfyui_exc_to_http(exc) from exc
    return {"ok": True, "job": job.to_dict()}


@app.post("/api/comfyui/models/download-file")
def api_comfyui_models_download_file(
    body: ComfyUIFileDownloadBody,
    _: None = Depends(require_mutation_auth),
) -> dict[str, Any]:
    mgr = get_download_manager()
    try:
        job = mgr.start_file_download(
            body.url.strip(),
            body.relative_path.strip(),
        )
    except ComfyUIModelError as exc:
        raise _comfyui_exc_to_http(exc) from exc
    return {"ok": True, "job": job.to_dict()}


@app.get("/api/comfyui/models/jobs/{job_id}")
def api_comfyui_models_job(
    job_id: str,
    _: None = Depends(require_comfyui_auth),
) -> dict[str, Any]:
    mgr = get_download_manager()
    try:
        job = mgr.get_job(job_id)
    except ComfyUIModelError as exc:
        raise _comfyui_exc_to_http(exc) from exc
    return {"job": job.to_dict()}


@app.delete("/api/comfyui/models/jobs/{job_id}")
def api_comfyui_models_job_cancel(
    job_id: str,
    cleanup: bool = Query(default=True),
    _: None = Depends(require_mutation_auth),
) -> dict[str, Any]:
    mgr = get_download_manager()
    try:
        job = mgr.cancel_job(job_id, cleanup=cleanup)
    except ComfyUIModelError as exc:
        raise _comfyui_exc_to_http(exc) from exc
    return {"ok": True, "job": job.to_dict()}


@app.delete("/api/comfyui/models/file")
def api_comfyui_models_delete_file(
    relative_path: str = Query(min_length=1),
    _: None = Depends(require_mutation_auth),
) -> dict[str, Any]:
    mgr = get_download_manager()
    try:
        result = mgr.delete_model_file(relative_path.strip())
    except ComfyUIModelError as exc:
        raise _comfyui_exc_to_http(exc) from exc
    return {"ok": True, **result}


@app.delete("/api/comfyui/models/{model_id}")
def api_comfyui_models_delete(
    model_id: str,
    _: None = Depends(require_mutation_auth),
) -> dict[str, Any]:
    mgr = get_download_manager()
    try:
        result = mgr.delete_model(model_id.strip())
    except ComfyUIModelError as exc:
        raise _comfyui_exc_to_http(exc) from exc
    return {"ok": True, **result}


@app.get("/api/files")
def api_files_list(
    path: str = Query(default="/"),
    _: None = Depends(require_file_auth),
) -> dict[str, Any]:
    try:
        return file_explorer.list_directory(path)
    except FileExplorerError as exc:
        raise _file_exc_to_http(exc) from exc


@app.get("/api/files/read")
def api_files_read(
    path: str = Query(...),
    _: None = Depends(require_file_auth),
) -> dict[str, Any]:
    try:
        return file_explorer.read_file(path)
    except FileExplorerError as exc:
        raise _file_exc_to_http(exc) from exc


@app.put("/api/files/write")
def api_files_write(
    body: FileWriteBody,
    _: None = Depends(require_file_auth),
) -> dict[str, Any]:
    target = file_explorer.resolve_path(body.path)
    if target.exists() and target.is_file():
        try:
            info = file_explorer.read_file(body.path)
            if info.get("binary"):
                raise HTTPException(
                    status_code=415,
                    detail="Nie można zapisać pliku binarnego jako tekst",
                )
        except FileExplorerError as exc:
            raise _file_exc_to_http(exc) from exc
    try:
        return file_explorer.write_file(body.path, body.content)
    except FileExplorerError as exc:
        raise _file_exc_to_http(exc) from exc


@app.post("/api/files/mkdir")
def api_files_mkdir(
    body: FilePathBody,
    _: None = Depends(require_file_auth),
) -> dict[str, Any]:
    try:
        return file_explorer.mkdir(body.path)
    except FileExplorerError as exc:
        raise _file_exc_to_http(exc) from exc


@app.post("/api/files/rename")
def api_files_rename(
    body: FileRenameBody,
    _: None = Depends(require_file_auth),
) -> dict[str, Any]:
    try:
        return file_explorer.rename_path(body.old_path, body.new_path)
    except FileExplorerError as exc:
        raise _file_exc_to_http(exc) from exc


@app.delete("/api/files")
def api_files_delete(
    path: str = Query(...),
    _: None = Depends(require_file_auth),
) -> dict[str, Any]:
    try:
        return file_explorer.delete_path(path)
    except FileExplorerError as exc:
        raise _file_exc_to_http(exc) from exc


@app.websocket("/api/cli/ws")
async def api_cli_ws(websocket: WebSocket) -> None:
    if not _cli_ws_auth_ok(websocket):
        await websocket.close(code=1008, reason="Invalid or missing API key")
        return

    if not CLI_ENABLED:
        await websocket.close(code=1008, reason="CLI disabled")
        return

    session_id = websocket.query_params.get("session_id", "").strip() or None

    if session_id:
        if get_session_manager().get(session_id) is None:
            await websocket.close(code=1008, reason="Session expired")
            return
    elif not sessions_available(CLI_MAX_SESSIONS):
        await websocket.close(
            code=1008,
            reason=f"Limit sesji CLI ({CLI_MAX_SESSIONS}) — spróbuj później",
        )
        return

    await websocket.accept()

    try:
        await run_pty_session(
            websocket,
            shell=CLI_SHELL,
            cwd=CLI_DEFAULT_CWD,
            env=os.environ.copy(),
            max_sessions=CLI_MAX_SESSIONS,
            session_id=session_id,
            session_ttl=CLI_SESSION_TTL,
        )
    except CliSessionLimitError:
        await websocket.close(code=1008, reason="Session limit")


def _validate_stacks(stacks: list[StackConfig]) -> None:
    for stack in stacks:
        compose_dir = stack.path
        if find_compose_file(compose_dir) is None:
            log.error(
                "Stack %s: brak compose w %s (REPO_ROOT=%s). "
                "Ustaw REPO_ROOT w /opt/control-plane/.env i zrestartuj server-ui.",
                stack.id,
                compose_dir,
                REPO_ROOT,
            )
            sys.exit(1)


def main() -> None:
    if HOST.strip() not in ("127.0.0.1", "localhost", "::1") and not API_KEY:
        log.error(
            "API_KEY is required when SERVER_UI_HOST=%s (LAN bind). Set API_KEY in .env",
            HOST,
        )
        sys.exit(1)

    if not STACKS:
        log.error("No stacks loaded from %s", CONFIG_PATH)
        sys.exit(1)

    _validate_stacks(STACKS)

    log.info("Repo root: %s", REPO_ROOT)
    if STACKS:
        log.info("First stack compose_dir: %s", STACKS[0].compose_dir)
    log.info("Stacks: %s", ", ".join(s.id for s in STACKS))
    log.info("GPU fan agent: %s", GPU_FAN_AGENT_URL)
    if API_KEY:
        src = api_key_source(STACK_DIR, ENV)
        log.info("API key: configured (source: %s)", src)
    else:
        log.warning("API_KEY not set (proxy may fail agent auth on LAN)")
    if CLI_ENABLED:
        log.info(
            "CLI: enabled (shell=%s, cwd=%s, max_sessions=%d, session_ttl=%ds)",
            CLI_SHELL,
            CLI_DEFAULT_CWD,
            CLI_MAX_SESSIONS,
            CLI_SESSION_TTL,
        )
    else:
        log.info("CLI: disabled (CLI_ENABLED=0)")
    log.info("Web UI at http://%s:%d", HOST if HOST != "0.0.0.0" else "0.0.0.0", PORT)
    uvicorn.run(app, host=HOST, port=PORT, log_level="info")


if __name__ == "__main__":
    main()