feat: add workflow studio and postgres support
This commit is contained in:
boos
@@ -1,4 +0,0 @@
|
||||
[submodule "comfy_source"]
|
||||
path = comfy_source
|
||||
url = https://github.com/comfyanonymous/ComfyUI.git
|
||||
branch = master
|
||||
+13
-14
@@ -4,15 +4,6 @@ version = "1.3.0"
|
||||
edition = "2021"
|
||||
description = "AI API Gateway — Rust rewrite of BoosAPI"
|
||||
|
||||
[features]
|
||||
# Default off so stock builds keep using the legacy embedded ComfyUI
|
||||
# singleton (`state.comfyui`). Enable once the per-user pool has been
|
||||
# validated in staging (see openspec/changes/comfyui-multi-tenant-api-only,
|
||||
# rollout step 9.2). Once 9.4 lands, this flag and the legacy path
|
||||
# disappear together.
|
||||
default = []
|
||||
multi_tenant_comfy = []
|
||||
|
||||
[dependencies]
|
||||
# Web framework
|
||||
axum = { version = "0.8", features = ["macros", "ws", "multipart"] }
|
||||
@@ -27,6 +18,8 @@ tokio = { version = "1", features = ["full"] }
|
||||
sqlx = { version = "0.8", features = [
|
||||
"runtime-tokio",
|
||||
"sqlite",
|
||||
"postgres",
|
||||
"any",
|
||||
"migrate",
|
||||
"chrono",
|
||||
] }
|
||||
@@ -75,11 +68,8 @@ once_cell = "1"
|
||||
parking_lot = "0.12"
|
||||
mime = "0.3"
|
||||
base64 = "0.22"
|
||||
|
||||
# WebSocket client (used by /comfy/* WS proxy to forward upstream to per-user containers)
|
||||
tokio-tungstenite = { version = "0.24", default-features = false, features = ["connect"] }
|
||||
async-trait = "0.1"
|
||||
urlencoding = "2"
|
||||
include_dir = "0.7"
|
||||
mime_guess = "2"
|
||||
|
||||
[dev-dependencies]
|
||||
rstest = "0.24"
|
||||
@@ -93,6 +83,15 @@ opt-level = 3
|
||||
lto = true
|
||||
codegen-units = 1
|
||||
|
||||
# Faster production-like build for staging/manual verification. This keeps
|
||||
# optimizations on but avoids the slow LTO + single-codegen-unit release path.
|
||||
[profile.fast-release]
|
||||
inherits = "release"
|
||||
opt-level = 2
|
||||
lto = false
|
||||
codegen-units = 256
|
||||
incremental = true
|
||||
|
||||
# Dev-profile speedups: line-tables-only debuginfo (much smaller .debug
|
||||
# sections → much faster link) and 256 codegen units (more parallelism
|
||||
# during codegen, slightly worse runtime — fine for dev). Dependencies
|
||||
|
||||
@@ -2,14 +2,14 @@
|
||||
|
||||
# BoosAPI
|
||||
|
||||
**AI 聚合网关 + ComfyUI Agent — 将分散的中转站整合为一个统一入口**
|
||||
**AI 聚合网关 + Workflow Studio — 将分散的中转站整合为一个统一入口**
|
||||
|
||||
<p>
|
||||
把你各处注册的 New API / One API / OneHub / DoneHub / Veloera / AnyRouter / Sub2API 等站点，
|
||||
<br>
|
||||
汇聚成 <strong>一个 API Key、一个入口</strong>，自动发现模型、智能路由、成本最优。
|
||||
<br>
|
||||
<strong>内置 ComfyUI 视频助手</strong> — 自然语言描述脚本，自动生成角色、场景、TTS、ComfyUI 工作流。
|
||||
<strong>Workflow Studio / 视频助手</strong> — 用 API-native 工作流和自然语言助手组合图像、视频、语音与文本能力。
|
||||
<br>
|
||||
<strong>SaaS 多租户架构</strong>，组织管理、成员邀请、套餐限制，开箱即用。
|
||||
</p>
|
||||
@@ -53,7 +53,7 @@
|
||||
| 📊 余额分散各处 | **集中看板** 一目了然，余额不足自动告警 |
|
||||
| ✅ 每天签到领额度 | **自动签到** 定时执行，奖励自动追踪 |
|
||||
| 🤷 不知道各站有什么模型 | **自动模型发现**，上游新增模型零配置 |
|
||||
| 🎨 想用 ComfyUI 但不想配工作流 | **对话助手**，自然语言描述即可生成图像/视频 |
|
||||
| 🎨 想编排图像/视频流程但不想跑本地模型 | **Workflow Studio + 对话助手**，直接复用 BoosAPI 模型路由 |
|
||||
| 👥 多人协作管理中转站 | **组织管理**，成员邀请，角色权限控制 |
|
||||
| 📋 不同团队需要不同配额 | **套餐系统**，Free/Pro/Enterprise 三级限制 |
|
||||
|
||||
@@ -61,13 +61,14 @@
|
||||
|
||||
## ✨ 核心功能
|
||||
|
||||
### 🎨 内置 ComfyUI 视频助手（Rust）
|
||||
### 🎨 Workflow Studio 与视频助手（Rust）
|
||||
|
||||
- **自然语言描述脚本**：说想法 → AI 自动提取角色、场景、对白
|
||||
- **自动生成四视图**：每个角色正/背/左/右 4 角度，调用 gpt-image-2 生成图片
|
||||
- **场景生成**：自动为每个场景生成配图
|
||||
- **TTS 语音合成**：对话台词自动转语音
|
||||
- **ComfyUI 工作流导出**：完整 workflow JSON，可直接导入 ComfyUI
|
||||
- **Workflow Studio**：API-native 工作流入口，不依赖 Docker、GPU 或本地 ComfyUI
|
||||
- **外部 ComfyUI 可选代理**：已有独立 ComfyUI 服务时，可通过 `comfyui.base_url` 接入 `/api/comfyui/proxy/*`
|
||||
- **Rust 后端**（boosapi-rs）：高性能 SSE 流式响应，零 Node.js 运行时依赖
|
||||
- 侧边栏「视频助手」入口，与聚合网关同一界面
|
||||
- **默认走本地模式**：复用本服务自身的 `/v1/*` 路由 + 启动时引导出来的内置 downstream key（`comfyui-agent`），开箱即用，无需额外配置
|
||||
@@ -209,8 +210,8 @@ docker compose up -d
|
||||
│ ├── db/ # 数据库模型与查询
|
||||
│ ├── routes/
|
||||
│ │ ├── mod.rs # 路由组装 + 静态文件 serve
|
||||
│ │ ├── comfyui_agent.rs # ComfyUI 视频助手 SSE 聊天
|
||||
│ │ ├── comfyui/ # ComfyUI 代理（prompt/history/view）
|
||||
│ │ ├── comfyui_agent.rs # 视频/工作流助手 SSE 聊天
|
||||
│ │ ├── comfyui/ # 外部 ComfyUI 代理（prompt/history/view）
|
||||
│ │ ├── proxy/ # OpenAI 兼容代理入口
|
||||
│ │ ├── auth.rs # 用户注册/登录 JWT
|
||||
│ │ ├── checkin.rs # 自动签到
|
||||
@@ -233,7 +234,8 @@ docker compose up -d
|
||||
│ ├── loginError.ts # 登录错误归一化（resolveLoginErrorMessage）
|
||||
│ ├── pages/
|
||||
│ │ ├── auth/ # 登录页拆分：LoginShell + AdminLoginPage + UserLoginPage（OAuthSlot 占位，行为由后续 change 接入）
|
||||
│ │ ├── ComfyUIAgent.tsx # 视频助手聊天界面
|
||||
│ │ ├── ComfyUIAgent.tsx # 视频/工作流助手聊天界面
|
||||
│ │ ├── Workflows.tsx # Workflow Studio 入口
|
||||
│ │ └── ...
|
||||
│ └── components/
|
||||
└── dist/web/ # 前端构建产物（由 Rust serve）
|
||||
@@ -254,10 +256,19 @@ BOOSAPI_AUTH__AUTH_TOKEN=your-admin-token
|
||||
BOOSAPI_AUTH__PROXY_TOKEN=sk-your-proxy-token
|
||||
BOOSAPI_AUTH__JWT_SECRET=your-jwt-secret
|
||||
|
||||
# 数据库
|
||||
# 数据库（默认 SQLite）
|
||||
BOOSAPI_DB__URL=sqlite:data/boosapi.db?mode=rwc
|
||||
|
||||
# 上游 API（ComfyUI Agent 回退用，默认本地模式时不需要设置）
|
||||
# PostgreSQL 准备/迁移
|
||||
# 当前版本保留 SQLite 作为运行主路径；可先把 SQLite 核心数据迁移到 PostgreSQL。
|
||||
# 完整 PostgreSQL 业务查询适配需要后续切换运行时 SQL 方言。
|
||||
# 示例：
|
||||
# boosapi migrate-sqlite-to-postgres \
|
||||
# --sqlite-url 'sqlite:data/boosapi.db?mode=rwc' \
|
||||
# --postgres-url 'postgres://boosapi:password@127.0.0.1:5432/boosapi'
|
||||
# 目标库已有数据时默认拒绝；确认覆盖时追加 --force-truncate。
|
||||
|
||||
# 上游 API（视频/工作流助手回退用，默认本地模式时不需要设置）
|
||||
# 设置 BOOSAPI_UPSTREAM_BASE 后切换到 fallback-remote 模式；BOOSAPI_UPSTREAM_KEY 缺省时仍用内置 comfyui-agent key
|
||||
BOOSAPI_UPSTREAM_BASE=https://api.boos.lat/v1
|
||||
BOOSAPI_UPSTREAM_KEY=sk-your-key
|
||||
@@ -272,12 +283,13 @@ BOOSAPI_WEB_DIR=dist/web
|
||||
|----|------|
|
||||
| 后端框架 | **Axum 0.8**（Rust） |
|
||||
| 运行时 | **Tokio**（异步运行时） |
|
||||
| 数据库 | **SQLite** + **SQLx**（编译期 SQL 检查） |
|
||||
| 数据库 | **SQLite** 运行主路径 + **PostgreSQL** schema/迁移准备（SQLx） |
|
||||
| 并行计算 | **Rayon**（CPU 密集任务） |
|
||||
| 认证鉴权 | **JWT** + API Key 双模式 |
|
||||
| 代理引擎 | 自定义 Token Router + Proxy Conductor |
|
||||
| 协议转换 | OpenAI / Claude / Gemini 全兼容 |
|
||||
| ComfyUI Agent | SSE 流式 + gpt-image-2 + TTS 自动生成 |
|
||||
| Workflow Studio | API-native 工作流入口，后续承接节点画布 |
|
||||
| 视频助手 | SSE 流式 + gpt-image-2 + TTS 自动生成 |
|
||||
| 前端框架 | React 19 + Vite（仅构建工具） |
|
||||
| 样式 | Tailwind CSS v4 |
|
||||
| 数据可视化 | VChart |
|
||||
|
||||
BIN
Binary file not shown.
+114
-80
@@ -6,48 +6,78 @@ set -euo pipefail
|
||||
ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
cd "$ROOT_DIR"
|
||||
|
||||
# ── 工具函数 ──────────────────────────────────────────────────────────────────
|
||||
|
||||
log() { echo "[$(date '+%H:%M:%S')] $*"; }
|
||||
ok() { echo "[$(date '+%H:%M:%S')] ✓ $*"; }
|
||||
fail() { echo "[$(date '+%H:%M:%S')] ✗ $*" >&2; exit 1; }
|
||||
|
||||
# ── 参数 ──────────────────────────────────────────────────────────────────────
|
||||
|
||||
RELEASE=0
|
||||
BUILD_PROFILE="${BOOSAPI_BUILD_PROFILE:-dev}"
|
||||
FRONTEND_ONLY=0
|
||||
BACKEND_ONLY=0
|
||||
NO_RUN=0
|
||||
REBUILD_IMAGE=0
|
||||
CARGO_FEATURES="${BOOSAPI_CARGO_FEATURES:-multi_tenant_comfy}"
|
||||
FORCE_FRONTEND=0
|
||||
SKIP_FRONTEND=0
|
||||
CARGO_FEATURES="${BOOSAPI_CARGO_FEATURES:-}"
|
||||
|
||||
usage() {
|
||||
cat <<'EOF'
|
||||
用法: ./build-and-run.sh [选项]
|
||||
|
||||
选项:
|
||||
--dev 使用 dev 模式编译后端（默认，最快）
|
||||
--release 使用 release 模式编译后端（更慢但性能好）
|
||||
--fast-release 使用 fast-release 模式编译后端（比 release 快，性能接近生产）
|
||||
--profile <name> 指定构建 profile：dev / release / fast-release
|
||||
--frontend-only 只构建前端
|
||||
--backend-only 只构建后端
|
||||
--force-frontend 强制重建前端，忽略缓存
|
||||
--skip-frontend 跳过前端构建（沿用现有 dist/web）
|
||||
--no-run 构建后不启动服务
|
||||
--rebuild-image 强制重新构建多租户 ComfyUI 镜像（默认：已存在则跳过）
|
||||
BOOSAPI_CARGO_FEATURES 可覆盖后端 Cargo features（默认：multi_tenant_comfy）
|
||||
BOOSAPI_BUILD_PROFILE 可设置默认 profile（dev / release / fast-release）
|
||||
BOOSAPI_CARGO_FEATURES 可覆盖后端 Cargo features（默认：空）
|
||||
-h, --help 显示帮助
|
||||
|
||||
示例:
|
||||
./build-and-run.sh --release
|
||||
./build-and-run.sh --fast-release
|
||||
BOOSAPI_BUILD_PROFILE=release ./build-and-run.sh
|
||||
EOF
|
||||
}
|
||||
|
||||
while [ $# -gt 0 ]; do
|
||||
case "$1" in
|
||||
--release) RELEASE=1 ;;
|
||||
--dev) BUILD_PROFILE="dev" ;;
|
||||
--release) BUILD_PROFILE="release" ;;
|
||||
--fast-release) BUILD_PROFILE="fast-release" ;;
|
||||
--profile)
|
||||
shift
|
||||
[ $# -gt 0 ] || fail "--profile 需要参数：dev / release / fast-release"
|
||||
BUILD_PROFILE="$1"
|
||||
;;
|
||||
--frontend-only) FRONTEND_ONLY=1 ;;
|
||||
--backend-only) BACKEND_ONLY=1 ;;
|
||||
--force-frontend) FORCE_FRONTEND=1 ;;
|
||||
--skip-frontend) SKIP_FRONTEND=1 ;;
|
||||
--no-run) NO_RUN=1 ;;
|
||||
--rebuild-image) REBUILD_IMAGE=1 ;;
|
||||
--rebuild-image) log "--rebuild-image 已废弃：内置 Docker ComfyUI 工作台已移除，忽略该参数" ;;
|
||||
-h|--help) usage; exit 0 ;;
|
||||
*) echo "未知参数: $1" >&2; usage >&2; exit 1 ;;
|
||||
esac
|
||||
shift
|
||||
done
|
||||
|
||||
# ── 工具函数 ──────────────────────────────────────────────────────────────────
|
||||
if [ "$FRONTEND_ONLY" -eq 1 ] && [ "$BACKEND_ONLY" -eq 1 ]; then
|
||||
fail "--frontend-only 和 --backend-only 不能同时使用"
|
||||
fi
|
||||
|
||||
log() { echo "[$(date '+%H:%M:%S')] $*"; }
|
||||
ok() { echo "[$(date '+%H:%M:%S')] ✓ $*"; }
|
||||
fail() { echo "[$(date '+%H:%M:%S')] ✗ $*" >&2; exit 1; }
|
||||
case "$BUILD_PROFILE" in
|
||||
dev|debug) BUILD_PROFILE="dev" ;;
|
||||
release|fast-release) ;;
|
||||
*) fail "未知构建 profile: $BUILD_PROFILE（可选：dev / release / fast-release）" ;;
|
||||
esac
|
||||
|
||||
server_port() {
|
||||
if [ -n "${BOOSAPI_SERVER__PORT:-}" ]; then
|
||||
@@ -115,7 +145,7 @@ stop_previous_server() {
|
||||
|
||||
while IFS= read -r pid; do
|
||||
[ -n "$pid" ] && binary_pids+=("$pid")
|
||||
done < <(pgrep -f "target/(debug|release)/boosapi" 2>/dev/null || true)
|
||||
done < <(pgrep -f "target/(debug|release|fast-release)/boosapi" 2>/dev/null || true)
|
||||
|
||||
stop_pids "boosapi binary" "${binary_pids[@]}"
|
||||
|
||||
@@ -131,40 +161,48 @@ stop_previous_server() {
|
||||
fi
|
||||
}
|
||||
|
||||
# ── 同步 ComfyUI 子模块 + 应用本地补丁 ─────────────────────────────────────────
|
||||
# comfy_source/ 是 vanilla ComfyUI 的 git 子模块，本地补丁存放在
|
||||
# comfy_patches/overlay/，由 comfy_patches/apply.sh 覆盖到子模块上。
|
||||
# Dockerfile 把整个 comfy_source/ 复制进镜像，所以这一步必须在 docker
|
||||
# build 之前完成，否则镜像里只会有 vanilla 代码。
|
||||
|
||||
if [ -f "$ROOT_DIR/.gitmodules" ] && [ -d "$ROOT_DIR/.git" ]; then
|
||||
if [ ! -f "$ROOT_DIR/comfy_source/main.py" ]; then
|
||||
log "初始化 ComfyUI 子模块..."
|
||||
git submodule update --init --recursive
|
||||
ok "子模块已初始化"
|
||||
build_frontend() {
|
||||
if [ "$SKIP_FRONTEND" -eq 1 ]; then
|
||||
log "跳过前端构建（--skip-frontend）"
|
||||
return 0
|
||||
fi
|
||||
if [ -d "$ROOT_DIR/comfy_patches/overlay" ]; then
|
||||
log "应用 ComfyUI 本地补丁 (comfy_patches/overlay → comfy_source/)..."
|
||||
"$ROOT_DIR/comfy_patches/apply.sh"
|
||||
ok "本地补丁已应用"
|
||||
fi
|
||||
fi
|
||||
|
||||
# ── 构建前端 ──────────────────────────────────────────────────────────────────
|
||||
|
||||
if [ "$BACKEND_ONLY" -eq 0 ]; then
|
||||
log "构建前端 (Vite)..."
|
||||
if ! command -v npx >/dev/null 2>&1; then
|
||||
fail "未找到 npx，请先安装 Node.js"
|
||||
fi
|
||||
|
||||
local fingerprint_file="$ROOT_DIR/dist/web/.boosapi-build-fingerprint"
|
||||
local fingerprint
|
||||
fingerprint="$(
|
||||
{
|
||||
printf 'build-script-v2\n'
|
||||
printf 'node='
|
||||
node --version 2>/dev/null || true
|
||||
printf 'npx='
|
||||
npx --version 2>/dev/null || true
|
||||
for path in package.json pnpm-lock.yaml package-lock.json yarn.lock vite.config.ts tsconfig.json tsconfig.web.json tsconfig.transpile.json src/web src/shared; do
|
||||
[ -e "$ROOT_DIR/$path" ] && find "$ROOT_DIR/$path" -type f -print
|
||||
done | sort | xargs -r sha256sum
|
||||
} | sha256sum | awk '{print $1}'
|
||||
)"
|
||||
|
||||
if [ "$FORCE_FRONTEND" -eq 0 ] \
|
||||
&& [ -f "$fingerprint_file" ] \
|
||||
&& [ -f "$ROOT_DIR/dist/web/index.html" ] \
|
||||
&& [ "$(cat "$fingerprint_file" 2>/dev/null || true)" = "$fingerprint" ]; then
|
||||
ok "前端无变化，跳过 Vite 构建 → dist/web/"
|
||||
return 0
|
||||
fi
|
||||
|
||||
log "构建前端 (Vite)..."
|
||||
# 直接调用 vite build，跳过 desktop:icons（脚本文件不存在）
|
||||
npx vite build
|
||||
mkdir -p "$ROOT_DIR/dist/web"
|
||||
printf '%s\n' "$fingerprint" > "$fingerprint_file"
|
||||
ok "前端构建完成 → dist/web/"
|
||||
fi
|
||||
}
|
||||
|
||||
# ── 构建后端 ──────────────────────────────────────────────────────────────────
|
||||
|
||||
if [ "$FRONTEND_ONLY" -eq 0 ]; then
|
||||
build_backend() {
|
||||
if ! command -v cargo >/dev/null 2>&1; then
|
||||
fail "未找到 cargo，请先安装 Rust 工具链"
|
||||
fi
|
||||
@@ -180,66 +218,62 @@ if [ "$FRONTEND_ONLY" -eq 0 ]; then
|
||||
CARGO_FEATURE_ARGS=(--features "$CARGO_FEATURES")
|
||||
fi
|
||||
|
||||
if [ "$RELEASE" -eq 1 ]; then
|
||||
if [ "$BUILD_PROFILE" = "release" ]; then
|
||||
log "构建后端 (cargo build --release ${CARGO_FEATURE_ARGS[*]})..."
|
||||
cargo build --release "${CARGO_FEATURE_ARGS[@]}"
|
||||
BINARY="$ROOT_DIR/target/release/boosapi"
|
||||
ok "后端构建完成 → target/release/boosapi"
|
||||
elif [ "$BUILD_PROFILE" = "fast-release" ]; then
|
||||
log "构建后端 (cargo build --profile fast-release ${CARGO_FEATURE_ARGS[*]})..."
|
||||
cargo build --profile fast-release "${CARGO_FEATURE_ARGS[@]}"
|
||||
BINARY="$ROOT_DIR/target/fast-release/boosapi"
|
||||
ok "后端构建完成 → target/fast-release/boosapi"
|
||||
else
|
||||
log "构建后端 (cargo build ${CARGO_FEATURE_ARGS[*]})..."
|
||||
cargo build "${CARGO_FEATURE_ARGS[@]}"
|
||||
BINARY="$ROOT_DIR/target/debug/boosapi"
|
||||
ok "后端构建完成 → target/debug/boosapi"
|
||||
fi
|
||||
fi
|
||||
}
|
||||
|
||||
# ── 构建多租户 ComfyUI 镜像 ───────────────────────────────────────────────────
|
||||
# 多租户 ComfyUI 镜像 (boosapi-comfyui:multi-tenant) 由 ComfyUIPool 在每次
|
||||
# 拉新用户容器时使用。镜像很重(~700MB),不必每次 ./build-and-run.sh 都
|
||||
# 重打。优先级:
|
||||
# 1. 已存在同名镜像 → 跳过(--rebuild-image 强制重建)
|
||||
# 2. 有 docker/boosapi-comfyui-multi-tenant-*.tar 备份 → docker load
|
||||
# 3. 都没有 → docker build
|
||||
#
|
||||
# `docker info` 用来确认 daemon 真的在跑,避免在只装了 client 的环境里
|
||||
# 卡 build。失败的话整段跳过,API 仍可走 legacy 单实例路径。
|
||||
|
||||
if [ "$FRONTEND_ONLY" -eq 0 ]; then
|
||||
if command -v docker >/dev/null 2>&1 && docker info >/dev/null 2>&1; then
|
||||
IMAGE="boosapi-comfyui:multi-tenant"
|
||||
SAVED_TAR="$(ls "$ROOT_DIR"/docker/boosapi-comfyui-multi-tenant-*.tar 2>/dev/null | head -1 || true)"
|
||||
DOCKERFILE="$ROOT_DIR/docker/Dockerfile.comfyui.multi-tenant"
|
||||
|
||||
if [ "$REBUILD_IMAGE" -eq 0 ] && docker image inspect "$IMAGE" >/dev/null 2>&1; then
|
||||
log "镜像 $IMAGE 已存在,跳过构建(--rebuild-image 强制重建)"
|
||||
elif [ "$REBUILD_IMAGE" -eq 0 ] && [ -n "$SAVED_TAR" ]; then
|
||||
log "从 $SAVED_TAR 加载 $IMAGE..."
|
||||
docker load -i "$SAVED_TAR"
|
||||
ok "镜像已加载 → $IMAGE"
|
||||
elif [ -f "$DOCKERFILE" ]; then
|
||||
log "构建多租户 ComfyUI 镜像 ($IMAGE)..."
|
||||
docker build \
|
||||
-f "$DOCKERFILE" \
|
||||
-t "$IMAGE" \
|
||||
"$ROOT_DIR"
|
||||
ok "ComfyUI 多租户镜像构建完成 → $IMAGE"
|
||||
else
|
||||
log "未找到 $DOCKERFILE,跳过镜像构建"
|
||||
cleanup_started_server() {
|
||||
if [ -n "${STARTED_PID:-}" ] && is_running_pid "$STARTED_PID"; then
|
||||
log "停止本次启动的后端进程: $STARTED_PID"
|
||||
kill -TERM "$STARTED_PID" 2>/dev/null || true
|
||||
wait "$STARTED_PID" 2>/dev/null || true
|
||||
fi
|
||||
else
|
||||
log "docker 不可用(未安装或 daemon 未运行),跳过多租户 ComfyUI 镜像准备"
|
||||
fi
|
||||
fi
|
||||
}
|
||||
|
||||
# ── 启动服务 ──────────────────────────────────────────────────────────────────
|
||||
# ── 构建 / 启动 ────────────────────────────────────────────────────────────────
|
||||
|
||||
if [ "$NO_RUN" -eq 1 ] || [ "$FRONTEND_ONLY" -eq 1 ]; then
|
||||
log "跳过启动（--no-run 或 --frontend-only）"
|
||||
if [ "$FRONTEND_ONLY" -eq 1 ]; then
|
||||
build_frontend
|
||||
log "跳过启动（--frontend-only）"
|
||||
exit 0
|
||||
fi
|
||||
|
||||
if [ "$BACKEND_ONLY" -ne 1 ]; then
|
||||
build_frontend
|
||||
fi
|
||||
|
||||
build_backend
|
||||
|
||||
if [ "$NO_RUN" -eq 1 ]; then
|
||||
log "跳过启动（--no-run）"
|
||||
exit 0
|
||||
fi
|
||||
|
||||
# 启动前清理旧后端进程和端口占用，避免 exec 后绑定端口失败。
|
||||
stop_previous_server
|
||||
|
||||
log "启动服务..."
|
||||
exec "$BINARY"
|
||||
if [ "$BACKEND_ONLY" -eq 1 ]; then
|
||||
exec "$BINARY"
|
||||
fi
|
||||
|
||||
"$BINARY" &
|
||||
STARTED_PID=$!
|
||||
trap cleanup_started_server INT TERM EXIT
|
||||
ok "后端已启动，PID: $STARTED_PID"
|
||||
|
||||
ok "前后端已就绪（profile: $BUILD_PROFILE），按 Ctrl+C 停止服务"
|
||||
wait "$STARTED_PID"
|
||||
|
||||
@@ -1,136 +0,0 @@
|
||||
# comfy_patches
|
||||
|
||||
Local patches that ride on top of vanilla ComfyUI (the `comfy_source/` git
|
||||
submodule).
|
||||
|
||||
`overlay/` mirrors the ComfyUI tree — every file under it is copied
|
||||
verbatim over the matching path inside `comfy_source/` by
|
||||
`apply.sh`. Run that whenever you bump the submodule or after a fresh
|
||||
clone with `--recurse-submodules`.
|
||||
|
||||
## Why an overlay instead of a fork
|
||||
|
||||
Pinning a fork would mean rebasing it every time we bump ComfyUI. An
|
||||
overlay keeps the upstream submodule clean: the patches are visible in
|
||||
git, easy to diff against vanilla, and trivial to drop once they're
|
||||
upstreamed.
|
||||
|
||||
## Patch inventory
|
||||
|
||||
| File | Type | Notes |
|
||||
|---|---|
|
||||
| `comfy_api_nodes/util/_helpers.py` | BoosAPI integration | `get_auth_header` falls back to `BOOSAPI_BUILTIN_KEY`; `default_base_url()` reads `--comfy-api-base`. Required for the multi-tenant pool — the per-user containers authenticate to BoosAPI through this. |
|
||||
| `app/frontend_management.py` + `server.py` + `openapi.yaml` | Feature | Adds `comfy_package_versions` to `/system_stats` so the admin UI can surface drift across `comfy*` packages. Candidate for upstreaming. |
|
||||
| `comfy_extras/nodes_audio.py` | Hardening | Guards Save/Preview/Trim/Split nodes against `audio is None` (happens when source video has no audio track). Candidate for upstreaming. |
|
||||
| `requirements.txt` | Bump | `comfyui-workflow-templates==0.9.77` (vanilla pins 0.9.75). |
|
||||
| `tests-unit/app_test/frontend_manager_test.py` | Test | Covers the new `get_comfy_package_versions()` path. |
|
||||
| `tests-unit/comfy_test/model_detection_test.py` | Test | Adjusted assertions for the audio-node guards. |
|
||||
| `README.md`, `SECURITY.md` | Docs | Local edits. Not load-bearing. |
|
||||
|
||||
## Apply
|
||||
|
||||
```bash
|
||||
./comfy_patches/apply.sh
|
||||
```
|
||||
|
||||
Idempotent — safe to re-run. Aborts if `comfy_source/` is dirty (so you
|
||||
don't silently clobber upstream work-in-progress).
|
||||
|
||||
## Upgrading vanilla ComfyUI
|
||||
|
||||
This is the SOP whenever you want to bump the `comfy_source/` submodule
|
||||
to a newer upstream commit / tag.
|
||||
|
||||
```bash
|
||||
cd /root/metapi
|
||||
|
||||
# 1. Pull the new upstream commit into the submodule.
|
||||
# Use --remote to follow the branch tracked in .gitmodules; pin to a
|
||||
# specific tag with `git -C comfy_source checkout <tag>` instead.
|
||||
git submodule update --remote comfy_source
|
||||
|
||||
# 2. Re-apply the local overlay on top.
|
||||
# apply.sh aborts if comfy_source/ is dirty — commit or stash first.
|
||||
./comfy_patches/apply.sh
|
||||
|
||||
# 3. Rebuild the multi-tenant image.
|
||||
# --no-run skips starting the boosapi binary; we only want the image.
|
||||
./build-and-run.sh --no-run
|
||||
|
||||
# 4. Smoke-test the new image without touching the live pool.
|
||||
docker run --rm -p 18188:8188 \
|
||||
-e BOOSAPI_BUILTIN_KEY=test \
|
||||
-e BOOSAPI_BASE=http://host.docker.internal:8787 \
|
||||
-e BOOSAPI_KEY=test \
|
||||
boosapi-comfyui:multi-tenant &
|
||||
sleep 5
|
||||
curl -s http://localhost:18188/system_stats | jq .comfy_package_versions
|
||||
# Expect a JSON object with comfy / comfyui_frontend_package versions.
|
||||
# Ctrl-C the container when done; --rm cleans it up.
|
||||
```
|
||||
|
||||
### When upstream changes a file we patch
|
||||
|
||||
If `apply.sh` finds that vanilla has changed under one of the overlay
|
||||
paths, it copies our overlay over anyway — vanilla's new content is
|
||||
silently lost. To detect this before bumping:
|
||||
|
||||
```bash
|
||||
# Diff our overlay against the new vanilla, file by file.
|
||||
for f in $(cd comfy_patches/overlay && find . -type f); do
|
||||
diff -u "comfy_source/$f" "comfy_patches/overlay/$f" \
|
||||
| head -40
|
||||
done
|
||||
```
|
||||
|
||||
For each diff that shows upstream behaviour we want to keep (not just
|
||||
formatting), edit `comfy_patches/overlay/<file>` to merge the
|
||||
intent — then re-run `apply.sh`. The patch-inventory table above is
|
||||
the canonical list of files to check.
|
||||
|
||||
### Backout
|
||||
|
||||
If the new upstream breaks something:
|
||||
|
||||
```bash
|
||||
# Roll the submodule back to whatever commit was pinned in HEAD.
|
||||
git submodule update comfy_source
|
||||
./comfy_patches/apply.sh
|
||||
./build-and-run.sh --no-run
|
||||
```
|
||||
|
||||
The previous image is still in `docker images` (we don't auto-prune),
|
||||
so you can also fall back to it via tag pin in production until the
|
||||
upgrade is fixed.
|
||||
|
||||
### Restoring from the saved image
|
||||
|
||||
A known-good build of the image is exported at
|
||||
`docker/boosapi-comfyui-multi-tenant-20260520.tar` (754 MiB). To
|
||||
restore on this or any other host:
|
||||
|
||||
```bash
|
||||
docker load -i docker/boosapi-comfyui-multi-tenant-20260520.tar
|
||||
# → boosapi-comfyui:multi-tenant
|
||||
```
|
||||
|
||||
Use this when a rebuild fails (network, dependency drift) and you need
|
||||
the running pool back fast — load the tar, the per-user containers
|
||||
will pick the image up on next launch.
|
||||
|
||||
## Runtime contract: ComfyUI calls our API, not the upstream's
|
||||
|
||||
The point of the `_helpers.py` overlay is to redirect every "call the
|
||||
ComfyUI cloud" code path inside the per-user container to BoosAPI:
|
||||
|
||||
- `--comfy-api-base http://host.docker.internal:8787/v1` is forwarded
|
||||
by `docker/comfyui-entrypoint.sh` from the `BOOSAPI_BASE` env var.
|
||||
- `BOOSAPI_BUILTIN_KEY` / `BOOSAPI_KEY` become the auth header that
|
||||
ComfyUI's API nodes send. BoosAPI verifies it, looks up the
|
||||
`token_routes` row for the requested model, picks a `route_channels`
|
||||
row, and forwards to the real upstream account.
|
||||
|
||||
So the per-user ComfyUI instances never talk to upstream platforms
|
||||
directly — every API node call funnels through BoosAPI's routing
|
||||
layer (the same path serving `/v1/chat/completions` for the rest of
|
||||
the product).
|
||||
@@ -1,49 +0,0 @@
|
||||
#!/usr/bin/env bash
|
||||
# Apply BoosAPI overlay patches on top of the vanilla ComfyUI submodule.
|
||||
#
|
||||
# Run after `git submodule update --init` or after bumping the submodule
|
||||
# pointer. Idempotent — files are overwritten with the overlay copy.
|
||||
#
|
||||
# Refuses to clobber a dirty submodule working tree so you don't lose
|
||||
# upstream work-in-progress.
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
|
||||
OVERLAY="$ROOT/comfy_patches/overlay"
|
||||
TARGET="$ROOT/comfy_source"
|
||||
|
||||
if [ ! -d "$TARGET" ]; then
|
||||
echo "comfy_source/ does not exist — did you forget --recurse-submodules?" >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
# Refuse to apply over a dirty submodule (anything other than our overlay diff).
|
||||
# The overlay files _will_ show as modified after we apply, so we only check
|
||||
# for OTHER changes by comparing against the upstream tree.
|
||||
if [ -d "$TARGET/.git" ] || [ -f "$TARGET/.git" ]; then
|
||||
cd "$TARGET"
|
||||
# Stash any pre-existing local changes that aren't part of the overlay.
|
||||
# We list overlay paths and exclude them from the dirty check.
|
||||
overlay_paths=$(cd "$OVERLAY" && find . -type f -printf '%P\n')
|
||||
unexpected=$(git status --porcelain | awk '{print $2}' | grep -vxF -f <(printf '%s\n' "$overlay_paths") || true)
|
||||
if [ -n "$unexpected" ]; then
|
||||
echo "comfy_source/ has unexpected local changes — aborting:" >&2
|
||||
echo "$unexpected" >&2
|
||||
echo "" >&2
|
||||
echo "Commit/stash inside the submodule, or rm -rf and re-init it." >&2
|
||||
exit 1
|
||||
fi
|
||||
cd "$ROOT"
|
||||
fi
|
||||
|
||||
count=0
|
||||
while IFS= read -r -d '' src; do
|
||||
rel="${src#"$OVERLAY/"}"
|
||||
dst="$TARGET/$rel"
|
||||
mkdir -p "$(dirname "$dst")"
|
||||
cp "$src" "$dst"
|
||||
count=$((count + 1))
|
||||
done < <(find "$OVERLAY" -type f -print0)
|
||||
|
||||
echo "Applied $count overlay file(s) to comfy_source/."
|
||||
@@ -1,479 +0,0 @@
|
||||
<div align="center">
|
||||
|
||||
# ComfyUI
|
||||
**The most powerful and modular AI engine for content creation.**
|
||||
|
||||
|
||||
[![Website][website-shield]][website-url]
|
||||
[![Dynamic JSON Badge][discord-shield]][discord-url]
|
||||
[![Twitter][twitter-shield]][twitter-url]
|
||||
[![Matrix][matrix-shield]][matrix-url]
|
||||
<br>
|
||||
[![][github-release-shield]][github-release-link]
|
||||
[![][github-release-date-shield]][github-release-link]
|
||||
[![][github-downloads-shield]][github-downloads-link]
|
||||
[![][github-downloads-latest-shield]][github-downloads-link]
|
||||
|
||||
[matrix-shield]: https://img.shields.io/badge/Matrix-000000?style=flat&logo=matrix&logoColor=white
|
||||
[matrix-url]: https://app.element.io/#/room/%23comfyui_space%3Amatrix.org
|
||||
[website-shield]: https://img.shields.io/badge/ComfyOrg-4285F4?style=flat
|
||||
[website-url]: https://www.comfy.org/
|
||||
<!-- Workaround to display total user from https://github.com/badges/shields/issues/4500#issuecomment-2060079995 -->
|
||||
[discord-shield]: https://img.shields.io/badge/dynamic/json?url=https%3A%2F%2Fdiscord.com%2Fapi%2Finvites%2Fcomfyorg%3Fwith_counts%3Dtrue&query=%24.approximate_member_count&logo=discord&logoColor=white&label=Discord&color=green&suffix=%20total
|
||||
[discord-url]: https://www.comfy.org/discord
|
||||
[twitter-shield]: https://img.shields.io/twitter/follow/ComfyUI
|
||||
[twitter-url]: https://x.com/ComfyUI
|
||||
|
||||
[github-release-shield]: https://img.shields.io/github/v/release/comfyanonymous/ComfyUI?style=flat&sort=semver
|
||||
[github-release-link]: https://github.com/comfyanonymous/ComfyUI/releases
|
||||
[github-release-date-shield]: https://img.shields.io/github/release-date/comfyanonymous/ComfyUI?style=flat
|
||||
[github-downloads-shield]: https://img.shields.io/github/downloads/comfyanonymous/ComfyUI/total?style=flat
|
||||
[github-downloads-latest-shield]: https://img.shields.io/github/downloads/comfyanonymous/ComfyUI/latest/total?style=flat&label=downloads%40latest
|
||||
[github-downloads-link]: https://github.com/comfyanonymous/ComfyUI/releases
|
||||
|
||||
<img width="1590" height="795" alt="ComfyUI Screenshot" src="https://github.com/user-attachments/assets/36e065e0-bfae-4456-8c7f-8369d5ea48a2" />
|
||||
<br>
|
||||
</div>
|
||||
|
||||
ComfyUI is the AI creation engine for visual professionals who demand control over every model, every parameter, and every output. Its powerful and modular node graph interface empowers creatives to generate images, videos, 3D models, audio, and more...
|
||||
- ComfyUI natively supports the latest open-source state of the art models.
|
||||
- API nodes provide access to the best closed source models such as Nano Banana, Seedance, Hunyuan3D, etc.
|
||||
- It is available on Windows, Linux, and macOS, locally with our desktop application or on our cloud.
|
||||
- The most sophisticated workflows can be exposed through a simple UI thanks to App Mode.
|
||||
- It integrates seamlessly into production pipelines with our API endpoints.
|
||||
|
||||
## Get Started
|
||||
|
||||
### Local
|
||||
|
||||
#### [Desktop Application](https://www.comfy.org/download)
|
||||
- The easiest way to get started.
|
||||
- Available on Windows & macOS.
|
||||
|
||||
#### [Windows Portable Package](#installing)
|
||||
- Get the latest commits and completely portable.
|
||||
- Available on Windows.
|
||||
|
||||
#### [Manual Install](#manual-install-windows-linux)
|
||||
Supports all operating systems and GPU types (NVIDIA, AMD, Intel, Apple Silicon, Ascend).
|
||||
|
||||
### Cloud
|
||||
|
||||
#### [Comfy Cloud](https://www.comfy.org/cloud)
|
||||
- Our official paid cloud version for those who can't afford local hardware.
|
||||
|
||||
## Examples
|
||||
See what ComfyUI can do with the [newer template workflows](https://comfy.org/workflows) or old [example workflows](https://comfyanonymous.github.io/ComfyUI_examples/).
|
||||
|
||||
## Features
|
||||
- Nodes/graph/flowchart interface to experiment and create complex Stable Diffusion workflows without needing to code anything.
|
||||
- NOTE: There are many more models supported than the list below, if you want to see what is supported see our templates list inside ComfyUI.
|
||||
- Image Models
|
||||
- SD1.x, SD2.x ([unCLIP](https://comfyanonymous.github.io/ComfyUI_examples/unclip/))
|
||||
- [SDXL](https://comfyanonymous.github.io/ComfyUI_examples/sdxl/), [SDXL Turbo](https://comfyanonymous.github.io/ComfyUI_examples/sdturbo/)
|
||||
- [Stable Cascade](https://comfyanonymous.github.io/ComfyUI_examples/stable_cascade/)
|
||||
- [SD3 and SD3.5](https://comfyanonymous.github.io/ComfyUI_examples/sd3/)
|
||||
- Pixart Alpha and Sigma
|
||||
- [AuraFlow](https://comfyanonymous.github.io/ComfyUI_examples/aura_flow/)
|
||||
- [HunyuanDiT](https://comfyanonymous.github.io/ComfyUI_examples/hunyuan_dit/)
|
||||
- [Flux](https://comfyanonymous.github.io/ComfyUI_examples/flux/)
|
||||
- [Lumina Image 2.0](https://comfyanonymous.github.io/ComfyUI_examples/lumina2/)
|
||||
- [HiDream](https://comfyanonymous.github.io/ComfyUI_examples/hidream/)
|
||||
- [Qwen Image](https://comfyanonymous.github.io/ComfyUI_examples/qwen_image/)
|
||||
- [Hunyuan Image 2.1](https://comfyanonymous.github.io/ComfyUI_examples/hunyuan_image/)
|
||||
- [Flux 2](https://comfyanonymous.github.io/ComfyUI_examples/flux2/)
|
||||
- [Z Image](https://comfyanonymous.github.io/ComfyUI_examples/z_image/)
|
||||
- Ernie Image
|
||||
- Image Editing Models
|
||||
- [Omnigen 2](https://comfyanonymous.github.io/ComfyUI_examples/omnigen/)
|
||||
- [Flux Kontext](https://comfyanonymous.github.io/ComfyUI_examples/flux/#flux-kontext-image-editing-model)
|
||||
- [HiDream E1.1](https://comfyanonymous.github.io/ComfyUI_examples/hidream/#hidream-e11)
|
||||
- [Qwen Image Edit](https://comfyanonymous.github.io/ComfyUI_examples/qwen_image/#edit-model)
|
||||
- Video Models
|
||||
- [Stable Video Diffusion](https://comfyanonymous.github.io/ComfyUI_examples/video/)
|
||||
- [Mochi](https://comfyanonymous.github.io/ComfyUI_examples/mochi/)
|
||||
- [LTX-Video](https://comfyanonymous.github.io/ComfyUI_examples/ltxv/)
|
||||
- [Hunyuan Video](https://comfyanonymous.github.io/ComfyUI_examples/hunyuan_video/)
|
||||
- [Wan 2.1](https://comfyanonymous.github.io/ComfyUI_examples/wan/)
|
||||
- [Wan 2.2](https://comfyanonymous.github.io/ComfyUI_examples/wan22/)
|
||||
- [Hunyuan Video 1.5](https://docs.comfy.org/tutorials/video/hunyuan/hunyuan-video-1-5)
|
||||
- Audio Models
|
||||
- [Stable Audio](https://comfyanonymous.github.io/ComfyUI_examples/audio/)
|
||||
- [ACE Step](https://comfyanonymous.github.io/ComfyUI_examples/audio/)
|
||||
- 3D Models
|
||||
- [Hunyuan3D 2.0](https://docs.comfy.org/tutorials/3d/hunyuan3D-2)
|
||||
- Asynchronous Queue system
|
||||
- Many optimizations: Only re-executes the parts of the workflow that changes between executions.
|
||||
- Smart memory management: can automatically run large models on GPUs with as low as 1GB vram with smart offloading.
|
||||
- Works even if you don't have a GPU with: ```--cpu``` (slow)
|
||||
- Can load ckpt and safetensors: All in one checkpoints or standalone diffusion models, VAEs and CLIP models.
|
||||
- Safe loading of ckpt, pt, pth, etc.. files.
|
||||
- Embeddings/Textual inversion
|
||||
- [Loras (regular, locon and loha)](https://comfyanonymous.github.io/ComfyUI_examples/lora/)
|
||||
- [Hypernetworks](https://comfyanonymous.github.io/ComfyUI_examples/hypernetworks/)
|
||||
- Loading full workflows (with seeds) from generated PNG, WebP and FLAC files.
|
||||
- Saving/Loading workflows as Json files.
|
||||
- Nodes interface can be used to create complex workflows like one for [Hires fix](https://comfyanonymous.github.io/ComfyUI_examples/2_pass_txt2img/) or much more advanced ones.
|
||||
- [Area Composition](https://comfyanonymous.github.io/ComfyUI_examples/area_composition/)
|
||||
- [Inpainting](https://comfyanonymous.github.io/ComfyUI_examples/inpaint/) with both regular and inpainting models.
|
||||
- [ControlNet and T2I-Adapter](https://comfyanonymous.github.io/ComfyUI_examples/controlnet/)
|
||||
- [Upscale Models (ESRGAN, ESRGAN variants, SwinIR, Swin2SR, etc...)](https://comfyanonymous.github.io/ComfyUI_examples/upscale_models/)
|
||||
- [GLIGEN](https://comfyanonymous.github.io/ComfyUI_examples/gligen/)
|
||||
- [Model Merging](https://comfyanonymous.github.io/ComfyUI_examples/model_merging/)
|
||||
- [LCM models and Loras](https://comfyanonymous.github.io/ComfyUI_examples/lcm/)
|
||||
- Latent previews with [TAESD](#how-to-show-high-quality-previews)
|
||||
- Works fully offline: core will never download anything unless you want to.
|
||||
- Optional API nodes to use paid models from external providers through the online [Comfy API](https://docs.comfy.org/tutorials/api-nodes/overview) disable with: `--disable-api-nodes`
|
||||
- [Config file](extra_model_paths.yaml.example) to set the search paths for models.
|
||||
|
||||
Workflow examples can be found on the [Examples page](https://comfyanonymous.github.io/ComfyUI_examples/)
|
||||
|
||||
## Release Process
|
||||
|
||||
ComfyUI follows a weekly release cycle targeting Monday but this regularly changes because of model releases or large changes to the codebase. There are three interconnected repositories:
|
||||
|
||||
1. **[ComfyUI Core](https://github.com/comfyanonymous/ComfyUI)**
|
||||
- Releases a new major stable version (e.g., v0.7.0) roughly every 2 weeks.
|
||||
- Starting from v0.4.0 patch versions will be used for fixes backported onto the current stable release.
|
||||
- Minor versions will be used for releases off the master branch.
|
||||
- Patch versions may still be used for releases on the master branch in cases where a backport would not make sense.
|
||||
- Commits outside of the stable release tags may be very unstable and break many custom nodes.
|
||||
- Serves as the foundation for the desktop release
|
||||
|
||||
2. **[ComfyUI Desktop](https://github.com/Comfy-Org/desktop)**
|
||||
- Builds a new release using the latest stable core version
|
||||
|
||||
3. **[ComfyUI Frontend](https://github.com/Comfy-Org/ComfyUI_frontend)**
|
||||
- Every 2+ weeks frontend updates are merged into the core repository
|
||||
- Features are frozen for the upcoming core release
|
||||
- Development continues for the next release cycle
|
||||
|
||||
## Shortcuts
|
||||
|
||||
| Keybind | Explanation |
|
||||
|------------------------------------|--------------------------------------------------------------------------------------------------------------------|
|
||||
| `Ctrl` + `Enter` | Queue up current graph for generation |
|
||||
| `Ctrl` + `Shift` + `Enter` | Queue up current graph as first for generation |
|
||||
| `Ctrl` + `Alt` + `Enter` | Cancel current generation |
|
||||
| `Ctrl` + `Z`/`Ctrl` + `Y` | Undo/Redo |
|
||||
| `Ctrl` + `S` | Save workflow |
|
||||
| `Ctrl` + `O` | Load workflow |
|
||||
| `Ctrl` + `A` | Select all nodes |
|
||||
| `Alt `+ `C` | Collapse/uncollapse selected nodes |
|
||||
| `Ctrl` + `M` | Mute/unmute selected nodes |
|
||||
| `Ctrl` + `B` | Bypass selected nodes (acts like the node was removed from the graph and the wires reconnected through) |
|
||||
| `Delete`/`Backspace` | Delete selected nodes |
|
||||
| `Ctrl` + `Backspace` | Delete the current graph |
|
||||
| `Space` | Move the canvas around when held and moving the cursor |
|
||||
| `Ctrl`/`Shift` + `Click` | Add clicked node to selection |
|
||||
| `Ctrl` + `C`/`Ctrl` + `V` | Copy and paste selected nodes (without maintaining connections to outputs of unselected nodes) |
|
||||
| `Ctrl` + `C`/`Ctrl` + `Shift` + `V` | Copy and paste selected nodes (maintaining connections from outputs of unselected nodes to inputs of pasted nodes) |
|
||||
| `Shift` + `Drag` | Move multiple selected nodes at the same time |
|
||||
| `Ctrl` + `D` | Load default graph |
|
||||
| `Alt` + `+` | Canvas Zoom in |
|
||||
| `Alt` + `-` | Canvas Zoom out |
|
||||
| `Ctrl` + `Shift` + LMB + Vertical drag | Canvas Zoom in/out |
|
||||
| `P` | Pin/Unpin selected nodes |
|
||||
| `Ctrl` + `G` | Group selected nodes |
|
||||
| `Q` | Toggle visibility of the queue |
|
||||
| `H` | Toggle visibility of history |
|
||||
| `R` | Refresh graph |
|
||||
| `F` | Show/Hide menu |
|
||||
| `.` | Fit view to selection (Whole graph when nothing is selected) |
|
||||
| Double-Click LMB | Open node quick search palette |
|
||||
| `Shift` + Drag | Move multiple wires at once |
|
||||
| `Ctrl` + `Alt` + LMB | Disconnect all wires from clicked slot |
|
||||
|
||||
`Ctrl` can also be replaced with `Cmd` instead for macOS users
|
||||
|
||||
# Installing
|
||||
|
||||
## Windows Portable
|
||||
|
||||
There is a portable standalone build for Windows that should work for running on Nvidia GPUs or for running on your CPU only on the [releases page](https://github.com/comfyanonymous/ComfyUI/releases).
|
||||
|
||||
### [Direct link to download](https://github.com/comfyanonymous/ComfyUI/releases/latest/download/ComfyUI_windows_portable_nvidia.7z)
|
||||
|
||||
Simply download, extract with [7-Zip](https://7-zip.org) or with the windows explorer on recent windows versions and run. For smaller models you normally only need to put the checkpoints (the huge ckpt/safetensors files) in: ComfyUI\models\checkpoints but many of the larger models have multiple files. Make sure to follow the instructions to know which subfolder to put them in ComfyUI\models\
|
||||
|
||||
If you have trouble extracting it, right click the file -> properties -> unblock
|
||||
|
||||
The portable above currently comes with python 3.13 and pytorch cuda 13.0. Update your Nvidia drivers if it doesn't start.
|
||||
|
||||
#### All Official Portable Downloads:
|
||||
|
||||
[Portable for AMD GPUs](https://github.com/comfyanonymous/ComfyUI/releases/latest/download/ComfyUI_windows_portable_amd.7z)
|
||||
|
||||
[Portable for Intel GPUs](https://github.com/comfyanonymous/ComfyUI/releases/latest/download/ComfyUI_windows_portable_intel.7z)
|
||||
|
||||
[Portable for Nvidia GPUs](https://github.com/comfyanonymous/ComfyUI/releases/latest/download/ComfyUI_windows_portable_nvidia.7z) (supports 20 series and above).
|
||||
|
||||
[Portable for Nvidia GPUs with pytorch cuda 12.6 and python 3.12](https://github.com/comfyanonymous/ComfyUI/releases/latest/download/ComfyUI_windows_portable_nvidia_cu126.7z) (Supports Nvidia 10 series and older GPUs).
|
||||
|
||||
#### How do I share models between another UI and ComfyUI?
|
||||
|
||||
See the [Config file](extra_model_paths.yaml.example) to set the search paths for models. In the standalone windows build you can find this file in the ComfyUI directory. Rename this file to extra_model_paths.yaml and edit it with your favorite text editor.
|
||||
|
||||
|
||||
## [comfy-cli](https://docs.comfy.org/comfy-cli/getting-started)
|
||||
|
||||
You can install and start ComfyUI using comfy-cli:
|
||||
```bash
|
||||
pip install comfy-cli
|
||||
comfy install
|
||||
```
|
||||
|
||||
## Manual Install (Windows, Linux)
|
||||
|
||||
Python 3.14 works but some custom nodes may have issues. The free threaded variant works but some dependencies will enable the GIL so it's not fully supported.
|
||||
|
||||
Python 3.13 is very well supported. If you have trouble with some custom node dependencies on 3.13 you can try 3.12
|
||||
|
||||
torch 2.4 and above is supported but some features and optimizations might only work on newer versions. We generally recommend using the latest major version of pytorch with the latest cuda version unless it is less than 2 weeks old.
|
||||
|
||||
### Instructions:
|
||||
|
||||
Git clone this repo.
|
||||
|
||||
Put your SD checkpoints (the huge ckpt/safetensors files) in: models/checkpoints
|
||||
|
||||
Put your VAE in: models/vae
|
||||
|
||||
|
||||
### AMD GPUs (Linux)
|
||||
|
||||
AMD users can install rocm and pytorch with pip if you don't have it already installed, this is the command to install the stable version:
|
||||
|
||||
```pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm7.2```
|
||||
|
||||
This is the command to install the nightly with ROCm 7.2 which might have some performance improvements:
|
||||
|
||||
```pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/rocm7.2```
|
||||
|
||||
|
||||
### AMD GPUs (Experimental: Windows and Linux), RDNA 3, 3.5 and 4 only.
|
||||
|
||||
These have less hardware support than the builds above but they work on windows. You also need to install the pytorch version specific to your hardware.
|
||||
|
||||
RDNA 3 (RX 7000 series):
|
||||
|
||||
```pip install --pre torch torchvision torchaudio --index-url https://rocm.nightlies.amd.com/v2/gfx110X-all/```
|
||||
|
||||
RDNA 3.5 (Strix halo/Ryzen AI Max+ 365):
|
||||
|
||||
```pip install --pre torch torchvision torchaudio --index-url https://rocm.nightlies.amd.com/v2/gfx1151/```
|
||||
|
||||
RDNA 4 (RX 9000 series):
|
||||
|
||||
```pip install --pre torch torchvision torchaudio --index-url https://rocm.nightlies.amd.com/v2/gfx120X-all/```
|
||||
|
||||
### Intel GPUs (Windows and Linux)
|
||||
|
||||
Intel Arc GPU users can install native PyTorch with torch.xpu support using pip. More information can be found [here](https://pytorch.org/docs/main/notes/get_start_xpu.html)
|
||||
|
||||
1. To install PyTorch xpu, use the following command:
|
||||
|
||||
```pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/xpu```
|
||||
|
||||
This is the command to install the Pytorch xpu nightly which might have some performance improvements:
|
||||
|
||||
```pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/xpu```
|
||||
|
||||
### NVIDIA
|
||||
|
||||
Nvidia users should install stable pytorch using this command:
|
||||
|
||||
```pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu130```
|
||||
|
||||
This is the command to install pytorch nightly instead which might have performance improvements.
|
||||
|
||||
```pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu132```
|
||||
|
||||
#### Troubleshooting
|
||||
|
||||
If you get the "Torch not compiled with CUDA enabled" error, uninstall torch with:
|
||||
|
||||
```pip uninstall torch```
|
||||
|
||||
And install it again with the command above.
|
||||
|
||||
### Dependencies
|
||||
|
||||
Install the dependencies by opening your terminal inside the ComfyUI folder and:
|
||||
|
||||
```pip install -r requirements.txt```
|
||||
|
||||
After this you should have everything installed and can proceed to running ComfyUI.
|
||||
|
||||
### Others:
|
||||
|
||||
#### Apple Mac silicon
|
||||
|
||||
You can install ComfyUI in Apple Mac silicon (M1 or M2) with any recent macOS version.
|
||||
|
||||
1. Install pytorch nightly. For instructions, read the [Accelerated PyTorch training on Mac](https://developer.apple.com/metal/pytorch/) Apple Developer guide (make sure to install the latest pytorch nightly).
|
||||
1. Follow the [ComfyUI manual installation](#manual-install-windows-linux) instructions for Windows and Linux.
|
||||
1. Install the ComfyUI [dependencies](#dependencies). If you have another Stable Diffusion UI [you might be able to reuse the dependencies](#i-already-have-another-ui-for-stable-diffusion-installed-do-i-really-have-to-install-all-of-these-dependencies).
|
||||
1. Launch ComfyUI by running `python main.py`
|
||||
|
||||
> **Note**: Remember to add your models, VAE, LoRAs etc. to the corresponding Comfy folders, as discussed in [ComfyUI manual installation](#manual-install-windows-linux).
|
||||
|
||||
#### Ascend NPUs
|
||||
|
||||
For models compatible with Ascend Extension for PyTorch (torch_npu). To get started, ensure your environment meets the prerequisites outlined on the [installation](https://ascend.github.io/docs/sources/ascend/quick_install.html) page. Here's a step-by-step guide tailored to your platform and installation method:
|
||||
|
||||
1. Begin by installing the recommended or newer kernel version for Linux as specified in the Installation page of torch-npu, if necessary.
|
||||
2. Proceed with the installation of Ascend Basekit, which includes the driver, firmware, and CANN, following the instructions provided for your specific platform.
|
||||
3. Next, install the necessary packages for torch-npu by adhering to the platform-specific instructions on the [Installation](https://ascend.github.io/docs/sources/pytorch/install.html#pytorch) page.
|
||||
4. Finally, adhere to the [ComfyUI manual installation](#manual-install-windows-linux) guide for Linux. Once all components are installed, you can run ComfyUI as described earlier.
|
||||
|
||||
#### Cambricon MLUs
|
||||
|
||||
For models compatible with Cambricon Extension for PyTorch (torch_mlu). Here's a step-by-step guide tailored to your platform and installation method:
|
||||
|
||||
1. Install the Cambricon CNToolkit by adhering to the platform-specific instructions on the [Installation](https://www.cambricon.com/docs/sdk_1.15.0/cntoolkit_3.7.2/cntoolkit_install_3.7.2/index.html)
|
||||
2. Next, install the PyTorch(torch_mlu) following the instructions on the [Installation](https://www.cambricon.com/docs/sdk_1.15.0/cambricon_pytorch_1.17.0/user_guide_1.9/index.html)
|
||||
3. Launch ComfyUI by running `python main.py`
|
||||
|
||||
#### Iluvatar Corex
|
||||
|
||||
For models compatible with Iluvatar Extension for PyTorch. Here's a step-by-step guide tailored to your platform and installation method:
|
||||
|
||||
1. Install the Iluvatar Corex Toolkit by adhering to the platform-specific instructions on the [Installation](https://support.iluvatar.com/#/DocumentCentre?id=1&nameCenter=2&productId=520117912052801536)
|
||||
2. Launch ComfyUI by running `python main.py`
|
||||
|
||||
|
||||
## [ComfyUI-Manager](https://github.com/Comfy-Org/ComfyUI-Manager/tree/manager-v4)
|
||||
|
||||
**ComfyUI-Manager** is an extension that allows you to easily install, update, and manage custom nodes for ComfyUI.
|
||||
|
||||
### Setup
|
||||
|
||||
1. Install the manager dependencies:
|
||||
```bash
|
||||
pip install -r manager_requirements.txt
|
||||
```
|
||||
|
||||
2. Enable the manager with the `--enable-manager` flag when running ComfyUI:
|
||||
```bash
|
||||
python main.py --enable-manager
|
||||
```
|
||||
|
||||
### Command Line Options
|
||||
|
||||
| Flag | Description |
|
||||
|------|-------------|
|
||||
| `--enable-manager` | Enable ComfyUI-Manager |
|
||||
| `--enable-manager-legacy-ui` | Use the legacy manager UI instead of the new UI (requires `--enable-manager`) |
|
||||
| `--disable-manager-ui` | Disable the manager UI and endpoints while keeping background features like security checks and scheduled installation completion (requires `--enable-manager`) |
|
||||
|
||||
|
||||
# Running
|
||||
|
||||
```python main.py```
|
||||
|
||||
### For AMD cards not officially supported by ROCm
|
||||
|
||||
Try running it with this command if you have issues:
|
||||
|
||||
For 6700, 6600 and maybe other RDNA2 or older: ```HSA_OVERRIDE_GFX_VERSION=10.3.0 python main.py```
|
||||
|
||||
For AMD 7600 and maybe other RDNA3 cards: ```HSA_OVERRIDE_GFX_VERSION=11.0.0 python main.py```
|
||||
|
||||
### AMD ROCm Tips
|
||||
|
||||
You can enable experimental memory efficient attention on recent pytorch in ComfyUI on some AMD GPUs using this command, it should already be enabled by default on RDNA3. If this improves speed for you on latest pytorch on your GPU please report it so that I can enable it by default.
|
||||
|
||||
```TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1 python main.py --use-pytorch-cross-attention```
|
||||
|
||||
You can also try setting this env variable `PYTORCH_TUNABLEOP_ENABLED=1` which might speed things up at the cost of a very slow initial run.
|
||||
|
||||
# Notes
|
||||
|
||||
Only parts of the graph that have an output with all the correct inputs will be executed.
|
||||
|
||||
Only parts of the graph that change from each execution to the next will be executed, if you submit the same graph twice only the first will be executed. If you change the last part of the graph only the part you changed and the part that depends on it will be executed.
|
||||
|
||||
Dragging a generated png on the webpage or loading one will give you the full workflow including seeds that were used to create it.
|
||||
|
||||
You can use () to change emphasis of a word or phrase like: (good code:1.2) or (bad code:0.8). The default emphasis for () is 1.1. To use () characters in your actual prompt escape them like \\( or \\).
|
||||
|
||||
You can use {day|night}, for wildcard/dynamic prompts. With this syntax "{wild|card|test}" will be randomly replaced by either "wild", "card" or "test" by the frontend every time you queue the prompt. To use {} characters in your actual prompt escape them like: \\{ or \\}.
|
||||
|
||||
Dynamic prompts also support C-style comments, like `// comment` or `/* comment */`.
|
||||
|
||||
To use a textual inversion concepts/embeddings in a text prompt put them in the models/embeddings directory and use them in the CLIPTextEncode node like this (you can omit the .pt extension):
|
||||
|
||||
```embedding:embedding_filename.pt```
|
||||
|
||||
|
||||
## How to show high-quality previews?
|
||||
|
||||
Use ```--preview-method auto``` to enable previews.
|
||||
|
||||
The default installation includes a fast latent preview method that's low-resolution. To enable higher-quality previews with [TAESD](https://github.com/madebyollin/taesd), download the [taesd_decoder.pth, taesdxl_decoder.pth, taesd3_decoder.pth and taef1_decoder.pth](https://github.com/madebyollin/taesd/) and place them in the `models/vae_approx` folder. Once they're installed, restart ComfyUI and launch it with `--preview-method taesd` to enable high-quality previews.
|
||||
|
||||
## How to use TLS/SSL?
|
||||
Generate a self-signed certificate (not appropriate for shared/production use) and key by running the command: `openssl req -x509 -newkey rsa:4096 -keyout key.pem -out cert.pem -sha256 -days 3650 -nodes -subj "/C=XX/ST=StateName/L=CityName/O=CompanyName/OU=CompanySectionName/CN=CommonNameOrHostname"`
|
||||
|
||||
Use `--tls-keyfile key.pem --tls-certfile cert.pem` to enable TLS/SSL, the app will now be accessible with `https://...` instead of `http://...`.
|
||||
|
||||
> Note: Windows users can use [alexisrolland/docker-openssl](https://github.com/alexisrolland/docker-openssl) or one of the [3rd party binary distributions](https://wiki.openssl.org/index.php/Binaries) to run the command example above.
|
||||
<br/><br/>If you use a container, note that the volume mount `-v` can be a relative path so `... -v ".\:/openssl-certs" ...` would create the key & cert files in the current directory of your command prompt or powershell terminal.
|
||||
|
||||
## Support and dev channel
|
||||
|
||||
[Discord](https://comfy.org/discord): Try the #help or #feedback channels.
|
||||
|
||||
[Matrix space: #comfyui_space:matrix.org](https://app.element.io/#/room/%23comfyui_space%3Amatrix.org) (it's like discord but open source).
|
||||
|
||||
See also: [https://www.comfy.org/](https://www.comfy.org/)
|
||||
|
||||
> _psst — we're hiring!_ Help build ComfyUI: [comfy.org/careers](https://www.comfy.org/careers)
|
||||
|
||||
## Frontend Development
|
||||
|
||||
As of August 15, 2024, we have transitioned to a new frontend, which is now hosted in a separate repository: [ComfyUI Frontend](https://github.com/Comfy-Org/ComfyUI_frontend). This repository now hosts the compiled JS (from TS/Vue) under the `web/` directory.
|
||||
|
||||
### Reporting Issues and Requesting Features
|
||||
|
||||
For any bugs, issues, or feature requests related to the frontend, please use the [ComfyUI Frontend repository](https://github.com/Comfy-Org/ComfyUI_frontend). This will help us manage and address frontend-specific concerns more efficiently.
|
||||
|
||||
### Using the Latest Frontend
|
||||
|
||||
The new frontend is now the default for ComfyUI. However, please note:
|
||||
|
||||
1. The frontend in the main ComfyUI repository is updated fortnightly.
|
||||
2. Daily releases are available in the separate frontend repository.
|
||||
|
||||
To use the most up-to-date frontend version:
|
||||
|
||||
1. For the latest daily release, launch ComfyUI with this command line argument:
|
||||
|
||||
```
|
||||
--front-end-version Comfy-Org/ComfyUI_frontend@latest
|
||||
```
|
||||
|
||||
2. For a specific version, replace `latest` with the desired version number:
|
||||
|
||||
```
|
||||
--front-end-version Comfy-Org/ComfyUI_frontend@1.2.2
|
||||
```
|
||||
|
||||
This approach allows you to easily switch between the stable fortnightly release and the cutting-edge daily updates, or even specific versions for testing purposes.
|
||||
|
||||
### Accessing the Legacy Frontend
|
||||
|
||||
If you need to use the legacy frontend for any reason, you can access it using the following command line argument:
|
||||
|
||||
```
|
||||
--front-end-version Comfy-Org/ComfyUI_legacy_frontend@latest
|
||||
```
|
||||
|
||||
This will use a snapshot of the legacy frontend preserved in the [ComfyUI Legacy Frontend repository](https://github.com/Comfy-Org/ComfyUI_legacy_frontend).
|
||||
|
||||
# QA
|
||||
|
||||
### Which GPU should I buy for this?
|
||||
|
||||
[See this page for some recommendations](https://github.com/comfyanonymous/ComfyUI/wiki/Which-GPU-should-I-buy-for-ComfyUI)
|
||||
@@ -1,44 +0,0 @@
|
||||
# Security Policy
|
||||
|
||||
## Scope
|
||||
|
||||
ComfyUI is designed to run locally. By default, the server binds to `127.0.0.1`, meaning only the user's own machine can reach it. Our threat model assumes:
|
||||
|
||||
- The user installed ComfyUI through a supported channel: the desktop application, the portable build, or a manual install following the README.
|
||||
- The user has not installed untrusted custom nodes. Custom nodes are arbitrary Python code and are trusted as much as any other software the user chooses to install.
|
||||
- Anyone with access to the ComfyUI URL is trusted (a direct consequence of the localhost-only default).
|
||||
- PyTorch and other dependencies are at the versions we ship or recommend in the README.
|
||||
|
||||
A report is in scope only if it affects a user operating within this threat model.
|
||||
|
||||
## What We Consider a Vulnerability
|
||||
|
||||
We want to hear about issues where a **reasonable user** — someone who does not install random untrusted nodes and who reads UI prompts and warnings before clicking through them — can be harmed by ComfyUI itself.
|
||||
|
||||
The clearest example: a workflow file that such a user might plausibly load and run, using only built-in nodes, that results in **untrusted code execution, arbitrary file read/write outside expected directories, or credential/data exfiltration**.
|
||||
|
||||
When submitting a report, please include a clear description of *why this is a problem for a typical local ComfyUI user*. Reports without this context are difficult to act on.
|
||||
|
||||
## What We Do Not Consider a Security Vulnerability
|
||||
|
||||
Please report the following through our regular [GitHub issues](https://github.com/comfyanonymous/ComfyUI/issues) instead. Filing them as security reports will likely cause them to be deprioritized or closed.
|
||||
|
||||
- **Issues requiring `--listen` or any non-default network exposure.** ComfyUI binds to localhost by default. If a remote attacker needs to reach the server for the attack to work, the user has chosen to expose it and is responsible for securing that deployment (firewall, reverse proxy, authentication, etc.). These are bugs, not vulnerabilities.
|
||||
- **`torch.load` and related deserialization issues in old PyTorch versions.** These are upstream PyTorch issues. Our distributions ship with — and our documentation recommends — recent PyTorch versions where these are addressed.
|
||||
- **Vulnerabilities that depend on outdated library versions** that we neither ship nor recommend (e.g., requiring PyTorch 2.6 or older).
|
||||
- **Issues that require a specific custom node to be installed.** Custom nodes are third-party code. Report these to the maintainer of that node.
|
||||
- **Crashes, hangs, or resource exhaustion from a loaded workflow.** Annoying, but not a security issue in our model. File a regular bug.
|
||||
- **Social-engineering scenarios** where the user is expected to ignore an explicit UI warning or prompt.
|
||||
|
||||
## Reporting
|
||||
|
||||
If you believe you have found an issue that falls within the scope above, please report it privately via GitHub's [Report a vulnerability](https://github.com/comfyanonymous/ComfyUI/security/advisories/new) feature rather than opening a public issue.
|
||||
|
||||
Please include:
|
||||
|
||||
1. A description of the vulnerability and the affected component.
|
||||
2. Reproduction steps, ideally with a minimal workflow file or proof-of-concept.
|
||||
3. The ComfyUI version, install method (desktop / portable / manual), and OS.
|
||||
4. An explanation of how this affects a typical local user as described in the threat model.
|
||||
|
||||
We will acknowledge valid reports and coordinate a fix and disclosure timeline with you.
|
||||
@@ -1,440 +0,0 @@
|
||||
from __future__ import annotations
|
||||
import argparse
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
import tempfile
|
||||
import zipfile
|
||||
import importlib
|
||||
from dataclasses import dataclass
|
||||
from functools import cached_property
|
||||
from pathlib import Path
|
||||
from typing import Dict, TypedDict, Optional
|
||||
from aiohttp import web
|
||||
from importlib.metadata import version
|
||||
|
||||
import requests
|
||||
from typing_extensions import NotRequired
|
||||
|
||||
from utils.install_util import get_missing_requirements_message, get_required_packages_versions
|
||||
|
||||
from comfy.cli_args import DEFAULT_VERSION_STRING
|
||||
import app.logger
|
||||
|
||||
|
||||
def frontend_install_warning_message():
|
||||
return f"""
|
||||
{get_missing_requirements_message()}
|
||||
|
||||
The ComfyUI frontend is shipped in a pip package so it needs to be updated separately from the ComfyUI code.
|
||||
""".strip()
|
||||
|
||||
def parse_version(version: str) -> tuple[int, int, int]:
|
||||
return tuple(map(int, version.split(".")))
|
||||
|
||||
def is_valid_version(version: str) -> bool:
|
||||
"""Validate if a string is a valid semantic version (X.Y.Z format)."""
|
||||
pattern = r"^(\d+)\.(\d+)\.(\d+)$"
|
||||
return bool(re.match(pattern, version))
|
||||
|
||||
def get_required_frontend_version():
|
||||
return get_required_packages_versions().get("comfyui-frontend-package", None)
|
||||
|
||||
|
||||
COMFY_PACKAGE_VERSIONS = []
|
||||
def get_comfy_package_versions():
|
||||
"""List installed/required versions for every comfy* package in requirements.txt."""
|
||||
if COMFY_PACKAGE_VERSIONS:
|
||||
return COMFY_PACKAGE_VERSIONS.copy()
|
||||
out = COMFY_PACKAGE_VERSIONS
|
||||
for name, required in (get_required_packages_versions() or {}).items():
|
||||
if not name.startswith("comfy"):
|
||||
continue
|
||||
try:
|
||||
installed = version(name)
|
||||
except Exception:
|
||||
installed = None
|
||||
out.append({"name": name, "installed": installed, "required": required})
|
||||
return out.copy()
|
||||
|
||||
|
||||
def check_comfy_packages_versions():
|
||||
"""Warn for every comfy* package whose installed version is below requirements.txt."""
|
||||
from packaging.version import InvalidVersion, parse as parse_pep440
|
||||
for pkg in get_comfy_package_versions():
|
||||
installed_str = pkg["installed"]
|
||||
required_str = pkg["required"]
|
||||
if not installed_str or not required_str:
|
||||
continue
|
||||
try:
|
||||
outdated = parse_pep440(installed_str) < parse_pep440(required_str)
|
||||
except InvalidVersion as e:
|
||||
logging.error(f"Failed to check {pkg['name']} version: {e}")
|
||||
continue
|
||||
if outdated:
|
||||
app.logger.log_startup_warning(
|
||||
f"""
|
||||
________________________________________________________________________
|
||||
WARNING WARNING WARNING WARNING WARNING
|
||||
|
||||
Installed {pkg["name"]} version {installed_str} is lower than the recommended version {required_str}.
|
||||
|
||||
{get_missing_requirements_message()}
|
||||
________________________________________________________________________
|
||||
""".strip()
|
||||
)
|
||||
else:
|
||||
logging.info("{} version: {}".format(pkg["name"], installed_str))
|
||||
|
||||
|
||||
REQUEST_TIMEOUT = 10 # seconds
|
||||
|
||||
|
||||
class Asset(TypedDict):
|
||||
url: str
|
||||
|
||||
|
||||
class Release(TypedDict):
|
||||
id: int
|
||||
tag_name: str
|
||||
name: str
|
||||
prerelease: bool
|
||||
created_at: str
|
||||
published_at: str
|
||||
body: str
|
||||
assets: NotRequired[list[Asset]]
|
||||
|
||||
|
||||
@dataclass
|
||||
class FrontEndProvider:
|
||||
owner: str
|
||||
repo: str
|
||||
|
||||
@property
|
||||
def folder_name(self) -> str:
|
||||
return f"{self.owner}_{self.repo}"
|
||||
|
||||
@property
|
||||
def release_url(self) -> str:
|
||||
return f"https://api.github.com/repos/{self.owner}/{self.repo}/releases"
|
||||
|
||||
@cached_property
|
||||
def all_releases(self) -> list[Release]:
|
||||
releases = []
|
||||
api_url = self.release_url
|
||||
while api_url:
|
||||
response = requests.get(api_url, timeout=REQUEST_TIMEOUT)
|
||||
response.raise_for_status() # Raises an HTTPError if the response was an error
|
||||
releases.extend(response.json())
|
||||
# GitHub uses the Link header to provide pagination links. Check if it exists and update api_url accordingly.
|
||||
if "next" in response.links:
|
||||
api_url = response.links["next"]["url"]
|
||||
else:
|
||||
api_url = None
|
||||
return releases
|
||||
|
||||
@cached_property
|
||||
def latest_release(self) -> Release:
|
||||
latest_release_url = f"{self.release_url}/latest"
|
||||
response = requests.get(latest_release_url, timeout=REQUEST_TIMEOUT)
|
||||
response.raise_for_status() # Raises an HTTPError if the response was an error
|
||||
return response.json()
|
||||
|
||||
@cached_property
|
||||
def latest_prerelease(self) -> Release:
|
||||
"""Get the latest pre-release version - even if it's older than the latest release"""
|
||||
release = [release for release in self.all_releases if release["prerelease"]]
|
||||
|
||||
if not release:
|
||||
raise ValueError("No pre-releases found")
|
||||
|
||||
# GitHub returns releases in reverse chronological order, so first is latest
|
||||
return release[0]
|
||||
|
||||
def get_release(self, version: str) -> Release:
|
||||
if version == "latest":
|
||||
return self.latest_release
|
||||
elif version == "prerelease":
|
||||
return self.latest_prerelease
|
||||
else:
|
||||
for release in self.all_releases:
|
||||
if release["tag_name"] in [version, f"v{version}"]:
|
||||
return release
|
||||
raise ValueError(f"Version {version} not found in releases")
|
||||
|
||||
|
||||
def download_release_asset_zip(release: Release, destination_path: str) -> None:
|
||||
"""Download dist.zip from github release."""
|
||||
asset_url = None
|
||||
for asset in release.get("assets", []):
|
||||
if asset["name"] == "dist.zip":
|
||||
asset_url = asset["url"]
|
||||
break
|
||||
|
||||
if not asset_url:
|
||||
raise ValueError("dist.zip not found in the release assets")
|
||||
|
||||
# Use a temporary file to download the zip content
|
||||
with tempfile.TemporaryFile() as tmp_file:
|
||||
headers = {"Accept": "application/octet-stream"}
|
||||
response = requests.get(
|
||||
asset_url, headers=headers, allow_redirects=True, timeout=REQUEST_TIMEOUT
|
||||
)
|
||||
response.raise_for_status() # Ensure we got a successful response
|
||||
|
||||
# Write the content to the temporary file
|
||||
tmp_file.write(response.content)
|
||||
|
||||
# Go back to the beginning of the temporary file
|
||||
tmp_file.seek(0)
|
||||
|
||||
# Extract the zip file content to the destination path
|
||||
with zipfile.ZipFile(tmp_file, "r") as zip_ref:
|
||||
zip_ref.extractall(destination_path)
|
||||
|
||||
|
||||
class FrontendManager:
|
||||
CUSTOM_FRONTENDS_ROOT = str(Path(__file__).parents[1] / "web_custom_versions")
|
||||
|
||||
@classmethod
|
||||
def get_required_frontend_version(cls) -> str:
|
||||
"""Get the required frontend package version."""
|
||||
return get_required_frontend_version()
|
||||
|
||||
@classmethod
|
||||
def get_installed_templates_version(cls) -> str:
|
||||
"""Get the currently installed workflow templates package version."""
|
||||
try:
|
||||
templates_version_str = version("comfyui-workflow-templates")
|
||||
return templates_version_str
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
@classmethod
|
||||
def get_required_templates_version(cls) -> str:
|
||||
return get_required_packages_versions().get("comfyui-workflow-templates", None)
|
||||
|
||||
@classmethod
|
||||
def get_comfy_package_versions(cls):
|
||||
"""List installed/required versions for every comfy* package in requirements.txt."""
|
||||
return get_comfy_package_versions()
|
||||
|
||||
@classmethod
|
||||
def default_frontend_path(cls) -> str:
|
||||
try:
|
||||
import comfyui_frontend_package
|
||||
|
||||
return str(importlib.resources.files(comfyui_frontend_package) / "static")
|
||||
except ImportError:
|
||||
logging.error(
|
||||
f"""
|
||||
********** ERROR ***********
|
||||
|
||||
comfyui-frontend-package is not installed.
|
||||
|
||||
{frontend_install_warning_message()}
|
||||
|
||||
********** ERROR ***********
|
||||
""".strip()
|
||||
)
|
||||
sys.exit(-1)
|
||||
|
||||
@classmethod
|
||||
def template_asset_map(cls) -> Optional[Dict[str, str]]:
|
||||
"""Return a mapping of template asset names to their absolute paths."""
|
||||
try:
|
||||
from comfyui_workflow_templates import (
|
||||
get_asset_path,
|
||||
iter_templates,
|
||||
)
|
||||
except ImportError:
|
||||
logging.error(
|
||||
f"""
|
||||
********** ERROR ***********
|
||||
|
||||
comfyui-workflow-templates is not installed.
|
||||
|
||||
{frontend_install_warning_message()}
|
||||
|
||||
********** ERROR ***********
|
||||
""".strip()
|
||||
)
|
||||
return None
|
||||
|
||||
try:
|
||||
template_entries = list(iter_templates())
|
||||
except Exception as exc:
|
||||
logging.error(f"Failed to enumerate workflow templates: {exc}")
|
||||
return None
|
||||
|
||||
asset_map: Dict[str, str] = {}
|
||||
try:
|
||||
for entry in template_entries:
|
||||
for asset in entry.assets:
|
||||
asset_map[asset.filename] = get_asset_path(
|
||||
entry.template_id, asset.filename
|
||||
)
|
||||
except Exception as exc:
|
||||
logging.error(f"Failed to resolve template asset paths: {exc}")
|
||||
return None
|
||||
|
||||
if not asset_map:
|
||||
logging.error("No workflow template assets found. Did the packages install correctly?")
|
||||
return None
|
||||
|
||||
return asset_map
|
||||
|
||||
|
||||
@classmethod
|
||||
def legacy_templates_path(cls) -> Optional[str]:
|
||||
"""Return the legacy templates directory shipped inside the meta package."""
|
||||
try:
|
||||
import comfyui_workflow_templates
|
||||
|
||||
return str(
|
||||
importlib.resources.files(comfyui_workflow_templates) / "templates"
|
||||
)
|
||||
except ImportError:
|
||||
logging.error(
|
||||
f"""
|
||||
********** ERROR ***********
|
||||
|
||||
comfyui-workflow-templates is not installed.
|
||||
|
||||
{frontend_install_warning_message()}
|
||||
|
||||
********** ERROR ***********
|
||||
""".strip()
|
||||
)
|
||||
return None
|
||||
|
||||
@classmethod
|
||||
def embedded_docs_path(cls) -> str:
|
||||
"""Get the path to embedded documentation"""
|
||||
try:
|
||||
import comfyui_embedded_docs
|
||||
|
||||
return str(
|
||||
importlib.resources.files(comfyui_embedded_docs) / "docs"
|
||||
)
|
||||
except ImportError:
|
||||
logging.info("comfyui-embedded-docs package not found")
|
||||
return None
|
||||
|
||||
@classmethod
|
||||
def parse_version_string(cls, value: str) -> tuple[str, str, str]:
|
||||
"""
|
||||
Args:
|
||||
value (str): The version string to parse.
|
||||
|
||||
Returns:
|
||||
tuple[str, str]: A tuple containing provider name and version.
|
||||
|
||||
Raises:
|
||||
argparse.ArgumentTypeError: If the version string is invalid.
|
||||
"""
|
||||
VERSION_PATTERN = r"^([a-zA-Z0-9][a-zA-Z0-9-]{0,38})/([a-zA-Z0-9_.-]+)@(v?\d+\.\d+\.\d+[-._a-zA-Z0-9]*|latest|prerelease)$"
|
||||
match_result = re.match(VERSION_PATTERN, value)
|
||||
if match_result is None:
|
||||
raise argparse.ArgumentTypeError(f"Invalid version string: {value}")
|
||||
|
||||
return match_result.group(1), match_result.group(2), match_result.group(3)
|
||||
|
||||
@classmethod
|
||||
def init_frontend_unsafe(
|
||||
cls, version_string: str, provider: Optional[FrontEndProvider] = None
|
||||
) -> str:
|
||||
"""
|
||||
Initializes the frontend for the specified version.
|
||||
|
||||
Args:
|
||||
version_string (str): The version string.
|
||||
provider (FrontEndProvider, optional): The provider to use. Defaults to None.
|
||||
|
||||
Returns:
|
||||
str: The path to the initialized frontend.
|
||||
|
||||
Raises:
|
||||
Exception: If there is an error during the initialization process.
|
||||
main error source might be request timeout or invalid URL.
|
||||
"""
|
||||
if version_string == DEFAULT_VERSION_STRING:
|
||||
check_comfy_packages_versions()
|
||||
return cls.default_frontend_path()
|
||||
|
||||
repo_owner, repo_name, version = cls.parse_version_string(version_string)
|
||||
|
||||
if version.startswith("v"):
|
||||
expected_path = str(
|
||||
Path(cls.CUSTOM_FRONTENDS_ROOT)
|
||||
/ f"{repo_owner}_{repo_name}"
|
||||
/ version.lstrip("v")
|
||||
)
|
||||
if os.path.exists(expected_path):
|
||||
logging.info(
|
||||
f"Using existing copy of specific frontend version tag: {repo_owner}/{repo_name}@{version}"
|
||||
)
|
||||
return expected_path
|
||||
|
||||
logging.info(
|
||||
f"Initializing frontend: {repo_owner}/{repo_name}@{version}, requesting version details from GitHub..."
|
||||
)
|
||||
|
||||
provider = provider or FrontEndProvider(repo_owner, repo_name)
|
||||
release = provider.get_release(version)
|
||||
|
||||
semantic_version = release["tag_name"].lstrip("v")
|
||||
web_root = str(
|
||||
Path(cls.CUSTOM_FRONTENDS_ROOT) / provider.folder_name / semantic_version
|
||||
)
|
||||
if not os.path.exists(web_root):
|
||||
try:
|
||||
os.makedirs(web_root, exist_ok=True)
|
||||
logging.info(
|
||||
"Downloading frontend(%s) version(%s) to (%s)",
|
||||
provider.folder_name,
|
||||
semantic_version,
|
||||
web_root,
|
||||
)
|
||||
logging.debug(release)
|
||||
download_release_asset_zip(release, destination_path=web_root)
|
||||
finally:
|
||||
# Clean up the directory if it is empty, i.e. the download failed
|
||||
if not os.listdir(web_root):
|
||||
os.rmdir(web_root)
|
||||
|
||||
return web_root
|
||||
|
||||
@classmethod
|
||||
def init_frontend(cls, version_string: str) -> str:
|
||||
"""
|
||||
Initializes the frontend with the specified version string.
|
||||
|
||||
Args:
|
||||
version_string (str): The version string to initialize the frontend with.
|
||||
|
||||
Returns:
|
||||
str: The path of the initialized frontend.
|
||||
"""
|
||||
try:
|
||||
return cls.init_frontend_unsafe(version_string)
|
||||
except Exception as e:
|
||||
logging.error("Failed to initialize frontend: %s", e)
|
||||
logging.info("Falling back to the default frontend.")
|
||||
check_comfy_packages_versions()
|
||||
return cls.default_frontend_path()
|
||||
@classmethod
|
||||
def template_asset_handler(cls):
|
||||
assets = cls.template_asset_map()
|
||||
if not assets:
|
||||
return None
|
||||
|
||||
async def serve_template(request: web.Request) -> web.StreamResponse:
|
||||
rel_path = request.match_info.get("path", "")
|
||||
target = assets.get(rel_path)
|
||||
if target is None:
|
||||
raise web.HTTPNotFound()
|
||||
return web.FileResponse(target)
|
||||
|
||||
return serve_template
|
||||
@@ -1,98 +0,0 @@
|
||||
import asyncio
|
||||
import contextlib
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
from collections.abc import Callable
|
||||
from io import BytesIO
|
||||
|
||||
from yarl import URL
|
||||
|
||||
from comfy.cli_args import args
|
||||
from comfy.model_management import processing_interrupted
|
||||
from comfy_api.latest import IO
|
||||
|
||||
from .common_exceptions import ProcessingInterrupted
|
||||
|
||||
_HAS_PCT_ESC = re.compile(r"%[0-9A-Fa-f]{2}") # any % followed by 2 hex digits
|
||||
_HAS_BAD_PCT = re.compile(r"%(?![0-9A-Fa-f]{2})") # any % not followed by 2 hex digits
|
||||
|
||||
|
||||
def is_processing_interrupted() -> bool:
|
||||
"""Return True if user/runtime requested interruption."""
|
||||
return processing_interrupted()
|
||||
|
||||
|
||||
def get_node_id(node_cls: type[IO.ComfyNode]) -> str:
|
||||
return node_cls.hidden.unique_id
|
||||
|
||||
|
||||
def get_auth_header(node_cls: type[IO.ComfyNode]) -> dict[str, str]:
|
||||
if node_cls.hidden.auth_token_comfy_org:
|
||||
return {"Authorization": f"Bearer {node_cls.hidden.auth_token_comfy_org}"}
|
||||
if node_cls.hidden.api_key_comfy_org:
|
||||
return {"X-API-KEY": node_cls.hidden.api_key_comfy_org}
|
||||
# BoosAPI fallback: when ComfyUI runs as an embedded subprocess of
|
||||
# BoosAPI, the parent injects the built-in `comfyui-agent` downstream
|
||||
# key via this env var so usage flows through the standard billing
|
||||
# pipeline. See src/services/comfyui_process.rs.
|
||||
builtin_key = os.environ.get("BOOSAPI_BUILTIN_KEY")
|
||||
if builtin_key:
|
||||
return {"Authorization": f"Bearer {builtin_key}"}
|
||||
return {}
|
||||
|
||||
|
||||
def default_base_url() -> str:
|
||||
return getattr(args, "comfy_api_base", "https://api.comfy.org")
|
||||
|
||||
|
||||
async def sleep_with_interrupt(
|
||||
seconds: float,
|
||||
node_cls: type[IO.ComfyNode] | None,
|
||||
label: str | None = None,
|
||||
start_ts: float | None = None,
|
||||
estimated_total: int | None = None,
|
||||
*,
|
||||
display_callback: Callable[[type[IO.ComfyNode], str, int, int | None], None] | None = None,
|
||||
):
|
||||
"""
|
||||
Sleep in 1s slices while:
|
||||
- Checking for interruption (raises ProcessingInterrupted).
|
||||
- Optionally emitting time progress via display_callback (if provided).
|
||||
"""
|
||||
end = time.monotonic() + seconds
|
||||
while True:
|
||||
if is_processing_interrupted():
|
||||
raise ProcessingInterrupted("Task cancelled")
|
||||
now = time.monotonic()
|
||||
if start_ts is not None and label and display_callback:
|
||||
with contextlib.suppress(Exception):
|
||||
display_callback(node_cls, label, int(now - start_ts), estimated_total)
|
||||
if now >= end:
|
||||
break
|
||||
await asyncio.sleep(min(1.0, end - now))
|
||||
|
||||
|
||||
def mimetype_to_extension(mime_type: str) -> str:
|
||||
"""Converts a MIME type to a file extension."""
|
||||
return mime_type.split("/")[-1].lower()
|
||||
|
||||
|
||||
def get_fs_object_size(path_or_object: str | BytesIO) -> int:
|
||||
if isinstance(path_or_object, str):
|
||||
return os.path.getsize(path_or_object)
|
||||
return len(path_or_object.getvalue())
|
||||
|
||||
|
||||
def to_aiohttp_url(url: str) -> URL:
|
||||
"""If `url` appears to be already percent-encoded (contains at least one valid %HH
|
||||
escape and no malformed '%' sequences) and contains no raw whitespace/control
|
||||
characters preserve the original encoding byte-for-byte (important for signed/presigned URLs).
|
||||
Otherwise, return `URL(url)` and allow yarl to normalize/quote as needed."""
|
||||
if any(c.isspace() for c in url) or any(ord(c) < 0x20 for c in url):
|
||||
# Avoid encoded=True if URL contains raw whitespace/control chars
|
||||
return URL(url)
|
||||
if _HAS_PCT_ESC.search(url) and not _HAS_BAD_PCT.search(url):
|
||||
# Preserve encoding only if it appears pre-encoded AND has no invalid % sequences
|
||||
return URL(url, encoded=True)
|
||||
return URL(url)
|
||||
@@ -1,841 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import av
|
||||
import torchaudio
|
||||
import torch
|
||||
import comfy.model_management
|
||||
import folder_paths
|
||||
import os
|
||||
import hashlib
|
||||
import node_helpers
|
||||
import logging
|
||||
from typing_extensions import override
|
||||
from comfy_api.latest import ComfyExtension, IO, UI
|
||||
|
||||
class EmptyLatentAudio(IO.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return IO.Schema(
|
||||
node_id="EmptyLatentAudio",
|
||||
display_name="Empty Latent Audio",
|
||||
category="latent/audio",
|
||||
essentials_category="Audio",
|
||||
inputs=[
|
||||
IO.Float.Input("seconds", default=47.6, min=1.0, max=1000.0, step=0.1),
|
||||
IO.Int.Input(
|
||||
"batch_size", default=1, min=1, max=4096, tooltip="The number of latent images in the batch.",
|
||||
),
|
||||
],
|
||||
outputs=[IO.Latent.Output()],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, seconds, batch_size) -> IO.NodeOutput:
|
||||
length = round((seconds * 44100 / 2048) / 2) * 2
|
||||
latent = torch.zeros([batch_size, 64, length], device=comfy.model_management.intermediate_device())
|
||||
return IO.NodeOutput({"samples":latent, "type": "audio"})
|
||||
|
||||
generate = execute # TODO: remove
|
||||
|
||||
|
||||
class ConditioningStableAudio(IO.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return IO.Schema(
|
||||
node_id="ConditioningStableAudio",
|
||||
category="conditioning",
|
||||
inputs=[
|
||||
IO.Conditioning.Input("positive"),
|
||||
IO.Conditioning.Input("negative"),
|
||||
IO.Float.Input("seconds_start", default=0.0, min=0.0, max=1000.0, step=0.1),
|
||||
IO.Float.Input("seconds_total", default=47.0, min=0.0, max=1000.0, step=0.1),
|
||||
],
|
||||
outputs=[
|
||||
IO.Conditioning.Output(display_name="positive"),
|
||||
IO.Conditioning.Output(display_name="negative"),
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, positive, negative, seconds_start, seconds_total) -> IO.NodeOutput:
|
||||
positive = node_helpers.conditioning_set_values(positive, {"seconds_start": seconds_start, "seconds_total": seconds_total})
|
||||
negative = node_helpers.conditioning_set_values(negative, {"seconds_start": seconds_start, "seconds_total": seconds_total})
|
||||
return IO.NodeOutput(positive, negative)
|
||||
|
||||
append = execute # TODO: remove
|
||||
|
||||
|
||||
class VAEEncodeAudio(IO.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return IO.Schema(
|
||||
node_id="VAEEncodeAudio",
|
||||
search_aliases=["audio to latent"],
|
||||
display_name="VAE Encode Audio",
|
||||
category="latent/audio",
|
||||
inputs=[
|
||||
IO.Audio.Input("audio"),
|
||||
IO.Vae.Input("vae"),
|
||||
],
|
||||
outputs=[IO.Latent.Output()],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, vae, audio) -> IO.NodeOutput:
|
||||
if audio is None:
|
||||
raise ValueError("VAEEncodeAudio: input audio is None (source video may have no audio track).")
|
||||
sample_rate = audio["sample_rate"]
|
||||
vae_sample_rate = getattr(vae, "audio_sample_rate", 44100)
|
||||
if vae_sample_rate != sample_rate:
|
||||
waveform = torchaudio.functional.resample(audio["waveform"], sample_rate, vae_sample_rate)
|
||||
else:
|
||||
waveform = audio["waveform"]
|
||||
|
||||
t = vae.encode(waveform.movedim(1, -1))
|
||||
return IO.NodeOutput({"samples": t})
|
||||
|
||||
encode = execute # TODO: remove
|
||||
|
||||
|
||||
def vae_decode_audio(vae, samples, tile=None, overlap=None):
|
||||
if tile is not None:
|
||||
audio = vae.decode_tiled(samples["samples"], tile_x=tile, tile_y=tile, overlap=overlap).movedim(-1, 1)
|
||||
else:
|
||||
audio = vae.decode(samples["samples"]).movedim(-1, 1)
|
||||
|
||||
std = torch.std(audio, dim=[1, 2], keepdim=True) * 5.0
|
||||
std[std < 1.0] = 1.0
|
||||
audio /= std
|
||||
vae_sample_rate = getattr(vae, "audio_sample_rate_output", getattr(vae, "audio_sample_rate", 44100))
|
||||
return {"waveform": audio, "sample_rate": vae_sample_rate if "sample_rate" not in samples else samples["sample_rate"]}
|
||||
|
||||
|
||||
class VAEDecodeAudio(IO.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return IO.Schema(
|
||||
node_id="VAEDecodeAudio",
|
||||
search_aliases=["latent to audio"],
|
||||
display_name="VAE Decode Audio",
|
||||
category="latent/audio",
|
||||
inputs=[
|
||||
IO.Latent.Input("samples"),
|
||||
IO.Vae.Input("vae"),
|
||||
],
|
||||
outputs=[IO.Audio.Output()],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, vae, samples) -> IO.NodeOutput:
|
||||
return IO.NodeOutput(vae_decode_audio(vae, samples))
|
||||
|
||||
decode = execute # TODO: remove
|
||||
|
||||
|
||||
class VAEDecodeAudioTiled(IO.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return IO.Schema(
|
||||
node_id="VAEDecodeAudioTiled",
|
||||
search_aliases=["latent to audio"],
|
||||
display_name="VAE Decode Audio (Tiled)",
|
||||
category="latent/audio",
|
||||
inputs=[
|
||||
IO.Latent.Input("samples"),
|
||||
IO.Vae.Input("vae"),
|
||||
IO.Int.Input("tile_size", default=512, min=32, max=8192, step=8),
|
||||
IO.Int.Input("overlap", default=64, min=0, max=1024, step=8),
|
||||
],
|
||||
outputs=[IO.Audio.Output()],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, vae, samples, tile_size, overlap) -> IO.NodeOutput:
|
||||
return IO.NodeOutput(vae_decode_audio(vae, samples, tile_size, overlap))
|
||||
|
||||
|
||||
class SaveAudio(IO.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return IO.Schema(
|
||||
node_id="SaveAudio",
|
||||
search_aliases=["export flac"],
|
||||
display_name="Save Audio (FLAC)",
|
||||
category="audio",
|
||||
essentials_category="Audio",
|
||||
inputs=[
|
||||
IO.Audio.Input("audio"),
|
||||
IO.String.Input("filename_prefix", default="audio/ComfyUI"),
|
||||
],
|
||||
hidden=[IO.Hidden.prompt, IO.Hidden.extra_pnginfo],
|
||||
is_output_node=True,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, audio, filename_prefix="ComfyUI", format="flac") -> IO.NodeOutput:
|
||||
if audio is None:
|
||||
raise ValueError("SaveAudio: input audio is None (source video may have no audio track).")
|
||||
return IO.NodeOutput(
|
||||
ui=UI.AudioSaveHelper.get_save_audio_ui(audio, filename_prefix=filename_prefix, cls=cls, format=format)
|
||||
)
|
||||
|
||||
save_flac = execute # TODO: remove
|
||||
|
||||
|
||||
class SaveAudioMP3(IO.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return IO.Schema(
|
||||
node_id="SaveAudioMP3",
|
||||
search_aliases=["export mp3"],
|
||||
display_name="Save Audio (MP3)",
|
||||
category="audio",
|
||||
essentials_category="Audio",
|
||||
inputs=[
|
||||
IO.Audio.Input("audio"),
|
||||
IO.String.Input("filename_prefix", default="audio/ComfyUI"),
|
||||
IO.Combo.Input("quality", options=["V0", "128k", "320k"], default="V0"),
|
||||
],
|
||||
hidden=[IO.Hidden.prompt, IO.Hidden.extra_pnginfo],
|
||||
is_output_node=True,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, audio, filename_prefix="ComfyUI", format="mp3", quality="128k") -> IO.NodeOutput:
|
||||
if audio is None:
|
||||
raise ValueError("SaveAudioMP3: input audio is None (source video may have no audio track).")
|
||||
return IO.NodeOutput(
|
||||
ui=UI.AudioSaveHelper.get_save_audio_ui(
|
||||
audio, filename_prefix=filename_prefix, cls=cls, format=format, quality=quality
|
||||
)
|
||||
)
|
||||
|
||||
save_mp3 = execute # TODO: remove
|
||||
|
||||
|
||||
class SaveAudioOpus(IO.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return IO.Schema(
|
||||
node_id="SaveAudioOpus",
|
||||
search_aliases=["export opus"],
|
||||
display_name="Save Audio (Opus)",
|
||||
category="audio",
|
||||
inputs=[
|
||||
IO.Audio.Input("audio"),
|
||||
IO.String.Input("filename_prefix", default="audio/ComfyUI"),
|
||||
IO.Combo.Input("quality", options=["64k", "96k", "128k", "192k", "320k"], default="128k"),
|
||||
],
|
||||
hidden=[IO.Hidden.prompt, IO.Hidden.extra_pnginfo],
|
||||
is_output_node=True,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, audio, filename_prefix="ComfyUI", format="opus", quality="V3") -> IO.NodeOutput:
|
||||
if audio is None:
|
||||
raise ValueError("SaveAudioOpus: input audio is None (source video may have no audio track).")
|
||||
return IO.NodeOutput(
|
||||
ui=UI.AudioSaveHelper.get_save_audio_ui(
|
||||
audio, filename_prefix=filename_prefix, cls=cls, format=format, quality=quality
|
||||
)
|
||||
)
|
||||
|
||||
save_opus = execute # TODO: remove
|
||||
|
||||
|
||||
class PreviewAudio(IO.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return IO.Schema(
|
||||
node_id="PreviewAudio",
|
||||
search_aliases=["play audio"],
|
||||
display_name="Preview Audio",
|
||||
category="audio",
|
||||
inputs=[
|
||||
IO.Audio.Input("audio"),
|
||||
],
|
||||
hidden=[IO.Hidden.prompt, IO.Hidden.extra_pnginfo],
|
||||
is_output_node=True,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, audio) -> IO.NodeOutput:
|
||||
if audio is None:
|
||||
raise ValueError("PreviewAudio: input audio is None (source video may have no audio track).")
|
||||
return IO.NodeOutput(ui=UI.PreviewAudio(audio, cls=cls))
|
||||
|
||||
save_flac = execute # TODO: remove
|
||||
|
||||
|
||||
def f32_pcm(wav: torch.Tensor) -> torch.Tensor:
|
||||
"""Convert audio to float 32 bits PCM format."""
|
||||
if wav.dtype.is_floating_point:
|
||||
return wav
|
||||
elif wav.dtype == torch.int16:
|
||||
return wav.float() / (2 ** 15)
|
||||
elif wav.dtype == torch.int32:
|
||||
return wav.float() / (2 ** 31)
|
||||
raise ValueError(f"Unsupported wav dtype: {wav.dtype}")
|
||||
|
||||
def load(filepath: str) -> tuple[torch.Tensor, int]:
|
||||
with av.open(filepath) as af:
|
||||
if not af.streams.audio:
|
||||
raise ValueError("No audio stream found in the file.")
|
||||
|
||||
stream = af.streams.audio[0]
|
||||
sr = stream.codec_context.sample_rate
|
||||
n_channels = stream.channels
|
||||
|
||||
frames = []
|
||||
length = 0
|
||||
for frame in af.decode(streams=stream.index):
|
||||
buf = torch.from_numpy(frame.to_ndarray())
|
||||
if buf.shape[0] != n_channels:
|
||||
buf = buf.view(-1, n_channels).t()
|
||||
|
||||
frames.append(buf)
|
||||
length += buf.shape[1]
|
||||
|
||||
if not frames:
|
||||
raise ValueError("No audio frames decoded.")
|
||||
|
||||
wav = torch.cat(frames, dim=1)
|
||||
wav = f32_pcm(wav)
|
||||
return wav, sr
|
||||
|
||||
class LoadAudio(IO.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
input_dir = folder_paths.get_input_directory()
|
||||
os.makedirs(input_dir, exist_ok=True)
|
||||
files = folder_paths.filter_files_content_types(os.listdir(input_dir), ["audio", "video"])
|
||||
return IO.Schema(
|
||||
node_id="LoadAudio",
|
||||
search_aliases=["import audio", "open audio", "audio file"],
|
||||
display_name="Load Audio",
|
||||
category="audio",
|
||||
essentials_category="Audio",
|
||||
inputs=[
|
||||
IO.Combo.Input("audio", upload=IO.UploadType.audio, options=sorted(files)),
|
||||
],
|
||||
outputs=[IO.Audio.Output()],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, audio) -> IO.NodeOutput:
|
||||
audio_path = folder_paths.get_annotated_filepath(audio)
|
||||
waveform, sample_rate = load(audio_path)
|
||||
audio = {"waveform": waveform.unsqueeze(0), "sample_rate": sample_rate}
|
||||
return IO.NodeOutput(audio)
|
||||
|
||||
@classmethod
|
||||
def fingerprint_inputs(cls, audio):
|
||||
image_path = folder_paths.get_annotated_filepath(audio)
|
||||
m = hashlib.sha256()
|
||||
with open(image_path, 'rb') as f:
|
||||
m.update(f.read())
|
||||
return m.digest().hex()
|
||||
|
||||
@classmethod
|
||||
def validate_inputs(cls, audio):
|
||||
if not folder_paths.exists_annotated_filepath(audio):
|
||||
return "Invalid audio file: {}".format(audio)
|
||||
return True
|
||||
|
||||
load = execute # TODO: remove
|
||||
|
||||
|
||||
class RecordAudio(IO.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return IO.Schema(
|
||||
node_id="RecordAudio",
|
||||
search_aliases=["microphone input", "audio capture", "voice input"],
|
||||
display_name="Record Audio",
|
||||
category="audio",
|
||||
inputs=[
|
||||
IO.Custom("AUDIO_RECORD").Input("audio"),
|
||||
],
|
||||
outputs=[IO.Audio.Output()],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, audio) -> IO.NodeOutput:
|
||||
audio_path = folder_paths.get_annotated_filepath(audio)
|
||||
|
||||
waveform, sample_rate = load(audio_path)
|
||||
audio = {"waveform": waveform.unsqueeze(0), "sample_rate": sample_rate}
|
||||
return IO.NodeOutput(audio)
|
||||
|
||||
load = execute # TODO: remove
|
||||
|
||||
|
||||
class TrimAudioDuration(IO.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return IO.Schema(
|
||||
node_id="TrimAudioDuration",
|
||||
search_aliases=["cut audio", "audio clip", "shorten audio"],
|
||||
display_name="Trim Audio Duration",
|
||||
description="Trim audio tensor into chosen time range.",
|
||||
category="audio",
|
||||
inputs=[
|
||||
IO.Audio.Input("audio"),
|
||||
IO.Float.Input(
|
||||
"start_index",
|
||||
default=0.0,
|
||||
min=-0xffffffffffffffff,
|
||||
max=0xffffffffffffffff,
|
||||
step=0.01,
|
||||
tooltip="Start time in seconds, can be negative to count from the end (supports sub-seconds).",
|
||||
),
|
||||
IO.Float.Input(
|
||||
"duration",
|
||||
default=60.0,
|
||||
min=0.0,
|
||||
step=0.01,
|
||||
tooltip="Duration in seconds",
|
||||
),
|
||||
],
|
||||
outputs=[IO.Audio.Output()],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, audio, start_index, duration) -> IO.NodeOutput:
|
||||
if audio is None:
|
||||
return IO.NodeOutput(None)
|
||||
waveform = audio["waveform"]
|
||||
sample_rate = audio["sample_rate"]
|
||||
audio_length = waveform.shape[-1]
|
||||
|
||||
if audio_length == 0:
|
||||
return IO.NodeOutput(audio)
|
||||
|
||||
if start_index < 0:
|
||||
start_frame = audio_length + int(round(start_index * sample_rate))
|
||||
else:
|
||||
start_frame = int(round(start_index * sample_rate))
|
||||
start_frame = max(0, min(start_frame, audio_length))
|
||||
|
||||
end_frame = start_frame + int(round(duration * sample_rate))
|
||||
end_frame = max(0, min(end_frame, audio_length))
|
||||
|
||||
if start_frame >= end_frame:
|
||||
raise ValueError("TrimAudioDuration: Start time must be less than end time and be within the audio length.")
|
||||
|
||||
return IO.NodeOutput({"waveform": waveform[..., start_frame:end_frame], "sample_rate": sample_rate})
|
||||
|
||||
trim = execute # TODO: remove
|
||||
|
||||
|
||||
class SplitAudioChannels(IO.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return IO.Schema(
|
||||
node_id="SplitAudioChannels",
|
||||
search_aliases=["stereo to mono"],
|
||||
display_name="Split Audio Channels",
|
||||
description="Separates the audio into left and right channels.",
|
||||
category="audio",
|
||||
inputs=[
|
||||
IO.Audio.Input("audio"),
|
||||
],
|
||||
outputs=[
|
||||
IO.Audio.Output(display_name="left"),
|
||||
IO.Audio.Output(display_name="right"),
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, audio) -> IO.NodeOutput:
|
||||
if audio is None:
|
||||
return IO.NodeOutput(None, None)
|
||||
waveform = audio["waveform"]
|
||||
sample_rate = audio["sample_rate"]
|
||||
|
||||
if waveform.shape[1] != 2:
|
||||
raise ValueError(f"AudioSplit: Input audio must be stereo (2 channels), got {waveform.shape[1]} channel(s).")
|
||||
|
||||
left_channel = waveform[..., 0:1, :]
|
||||
right_channel = waveform[..., 1:2, :]
|
||||
|
||||
return IO.NodeOutput({"waveform": left_channel, "sample_rate": sample_rate}, {"waveform": right_channel, "sample_rate": sample_rate})
|
||||
|
||||
separate = execute # TODO: remove
|
||||
|
||||
class JoinAudioChannels(IO.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return IO.Schema(
|
||||
node_id="JoinAudioChannels",
|
||||
display_name="Join Audio Channels",
|
||||
description="Joins left and right mono audio channels into a stereo audio.",
|
||||
category="audio",
|
||||
inputs=[
|
||||
IO.Audio.Input("audio_left"),
|
||||
IO.Audio.Input("audio_right"),
|
||||
],
|
||||
outputs=[
|
||||
IO.Audio.Output(display_name="audio"),
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, audio_left, audio_right) -> IO.NodeOutput:
|
||||
if audio_left is None and audio_right is None:
|
||||
return IO.NodeOutput(None)
|
||||
if audio_left is None:
|
||||
return IO.NodeOutput(audio_right)
|
||||
if audio_right is None:
|
||||
return IO.NodeOutput(audio_left)
|
||||
waveform_left = audio_left["waveform"]
|
||||
sample_rate_left = audio_left["sample_rate"]
|
||||
waveform_right = audio_right["waveform"]
|
||||
sample_rate_right = audio_right["sample_rate"]
|
||||
|
||||
if waveform_left.shape[1] != 1 or waveform_right.shape[1] != 1:
|
||||
raise ValueError("AudioJoin: Both input audios must be mono.")
|
||||
|
||||
# Handle different sample rates by resampling to the higher rate
|
||||
waveform_left, waveform_right, output_sample_rate = match_audio_sample_rates(
|
||||
waveform_left, sample_rate_left, waveform_right, sample_rate_right
|
||||
)
|
||||
|
||||
# Handle different lengths by trimming to the shorter length
|
||||
length_left = waveform_left.shape[-1]
|
||||
length_right = waveform_right.shape[-1]
|
||||
|
||||
if length_left != length_right:
|
||||
min_length = min(length_left, length_right)
|
||||
if length_left > min_length:
|
||||
logging.info(f"JoinAudioChannels: Trimming left channel from {length_left} to {min_length} samples.")
|
||||
waveform_left = waveform_left[..., :min_length]
|
||||
if length_right > min_length:
|
||||
logging.info(f"JoinAudioChannels: Trimming right channel from {length_right} to {min_length} samples.")
|
||||
waveform_right = waveform_right[..., :min_length]
|
||||
|
||||
# Join the channels into stereo
|
||||
left_channel = waveform_left[..., 0:1, :]
|
||||
right_channel = waveform_right[..., 0:1, :]
|
||||
stereo_waveform = torch.cat([left_channel, right_channel], dim=1)
|
||||
|
||||
return IO.NodeOutput({"waveform": stereo_waveform, "sample_rate": output_sample_rate})
|
||||
|
||||
|
||||
def match_audio_sample_rates(waveform_1, sample_rate_1, waveform_2, sample_rate_2):
|
||||
if sample_rate_1 != sample_rate_2:
|
||||
if sample_rate_1 > sample_rate_2:
|
||||
waveform_2 = torchaudio.functional.resample(waveform_2, sample_rate_2, sample_rate_1)
|
||||
output_sample_rate = sample_rate_1
|
||||
logging.info(f"Resampling audio2 from {sample_rate_2}Hz to {sample_rate_1}Hz for merging.")
|
||||
else:
|
||||
waveform_1 = torchaudio.functional.resample(waveform_1, sample_rate_1, sample_rate_2)
|
||||
output_sample_rate = sample_rate_2
|
||||
logging.info(f"Resampling audio1 from {sample_rate_1}Hz to {sample_rate_2}Hz for merging.")
|
||||
else:
|
||||
output_sample_rate = sample_rate_1
|
||||
return waveform_1, waveform_2, output_sample_rate
|
||||
|
||||
|
||||
class AudioConcat(IO.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return IO.Schema(
|
||||
node_id="AudioConcat",
|
||||
search_aliases=["join audio", "combine audio", "append audio"],
|
||||
display_name="Audio Concat",
|
||||
description="Concatenates the audio1 to audio2 in the specified direction.",
|
||||
category="audio",
|
||||
inputs=[
|
||||
IO.Audio.Input("audio1"),
|
||||
IO.Audio.Input("audio2"),
|
||||
IO.Combo.Input(
|
||||
"direction",
|
||||
options=['after', 'before'],
|
||||
default="after",
|
||||
tooltip="Whether to append audio2 after or before audio1.",
|
||||
)
|
||||
],
|
||||
outputs=[IO.Audio.Output()],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, audio1, audio2, direction) -> IO.NodeOutput:
|
||||
if audio1 is None and audio2 is None:
|
||||
return IO.NodeOutput(None)
|
||||
if audio1 is None:
|
||||
return IO.NodeOutput(audio2)
|
||||
if audio2 is None:
|
||||
return IO.NodeOutput(audio1)
|
||||
waveform_1 = audio1["waveform"]
|
||||
waveform_2 = audio2["waveform"]
|
||||
sample_rate_1 = audio1["sample_rate"]
|
||||
sample_rate_2 = audio2["sample_rate"]
|
||||
|
||||
if waveform_1.shape[1] == 1:
|
||||
waveform_1 = waveform_1.repeat(1, 2, 1)
|
||||
logging.info("AudioConcat: Converted mono audio1 to stereo by duplicating the channel.")
|
||||
if waveform_2.shape[1] == 1:
|
||||
waveform_2 = waveform_2.repeat(1, 2, 1)
|
||||
logging.info("AudioConcat: Converted mono audio2 to stereo by duplicating the channel.")
|
||||
|
||||
waveform_1, waveform_2, output_sample_rate = match_audio_sample_rates(waveform_1, sample_rate_1, waveform_2, sample_rate_2)
|
||||
|
||||
if direction == 'after':
|
||||
concatenated_audio = torch.cat((waveform_1, waveform_2), dim=2)
|
||||
elif direction == 'before':
|
||||
concatenated_audio = torch.cat((waveform_2, waveform_1), dim=2)
|
||||
|
||||
return IO.NodeOutput({"waveform": concatenated_audio, "sample_rate": output_sample_rate})
|
||||
|
||||
concat = execute # TODO: remove
|
||||
|
||||
|
||||
class AudioMerge(IO.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return IO.Schema(
|
||||
node_id="AudioMerge",
|
||||
search_aliases=["mix audio", "overlay audio", "layer audio"],
|
||||
display_name="Audio Merge",
|
||||
description="Combine two audio tracks by overlaying their waveforms.",
|
||||
category="audio",
|
||||
inputs=[
|
||||
IO.Audio.Input("audio1"),
|
||||
IO.Audio.Input("audio2"),
|
||||
IO.Combo.Input(
|
||||
"merge_method",
|
||||
options=["add", "mean", "subtract", "multiply"],
|
||||
tooltip="The method used to combine the audio waveforms.",
|
||||
)
|
||||
],
|
||||
outputs=[IO.Audio.Output()],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, audio1, audio2, merge_method) -> IO.NodeOutput:
|
||||
if audio1 is None and audio2 is None:
|
||||
return IO.NodeOutput(None)
|
||||
if audio1 is None:
|
||||
return IO.NodeOutput(audio2)
|
||||
if audio2 is None:
|
||||
return IO.NodeOutput(audio1)
|
||||
waveform_1 = audio1["waveform"]
|
||||
waveform_2 = audio2["waveform"]
|
||||
sample_rate_1 = audio1["sample_rate"]
|
||||
sample_rate_2 = audio2["sample_rate"]
|
||||
|
||||
waveform_1, waveform_2, output_sample_rate = match_audio_sample_rates(waveform_1, sample_rate_1, waveform_2, sample_rate_2)
|
||||
|
||||
length_1 = waveform_1.shape[-1]
|
||||
length_2 = waveform_2.shape[-1]
|
||||
|
||||
if length_1 == 0 or length_2 == 0:
|
||||
return IO.NodeOutput({"waveform": waveform_1, "sample_rate": output_sample_rate})
|
||||
|
||||
if length_2 > length_1:
|
||||
logging.info(f"AudioMerge: Trimming audio2 from {length_2} to {length_1} samples to match audio1 length.")
|
||||
waveform_2 = waveform_2[..., :length_1]
|
||||
elif length_2 < length_1:
|
||||
logging.info(f"AudioMerge: Padding audio2 from {length_2} to {length_1} samples to match audio1 length.")
|
||||
pad_shape = list(waveform_2.shape)
|
||||
pad_shape[-1] = length_1 - length_2
|
||||
pad_tensor = torch.zeros(pad_shape, dtype=waveform_2.dtype, device=waveform_2.device)
|
||||
waveform_2 = torch.cat((waveform_2, pad_tensor), dim=-1)
|
||||
|
||||
if merge_method == "add":
|
||||
waveform = waveform_1 + waveform_2
|
||||
elif merge_method == "subtract":
|
||||
waveform = waveform_1 - waveform_2
|
||||
elif merge_method == "multiply":
|
||||
waveform = waveform_1 * waveform_2
|
||||
elif merge_method == "mean":
|
||||
waveform = (waveform_1 + waveform_2) / 2
|
||||
|
||||
max_val = waveform.abs().max()
|
||||
if max_val > 1.0:
|
||||
waveform = waveform / max_val
|
||||
|
||||
return IO.NodeOutput({"waveform": waveform, "sample_rate": output_sample_rate})
|
||||
|
||||
merge = execute # TODO: remove
|
||||
|
||||
|
||||
class AudioAdjustVolume(IO.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return IO.Schema(
|
||||
node_id="AudioAdjustVolume",
|
||||
search_aliases=["audio gain", "loudness", "audio level"],
|
||||
display_name="Audio Adjust Volume",
|
||||
category="audio",
|
||||
inputs=[
|
||||
IO.Audio.Input("audio"),
|
||||
IO.Int.Input(
|
||||
"volume",
|
||||
default=1,
|
||||
min=-100,
|
||||
max=100,
|
||||
tooltip="Volume adjustment in decibels (dB). 0 = no change, +6 = double, -6 = half, etc",
|
||||
)
|
||||
],
|
||||
outputs=[IO.Audio.Output()],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, audio, volume) -> IO.NodeOutput:
|
||||
if audio is None:
|
||||
return IO.NodeOutput(None)
|
||||
if volume == 0:
|
||||
return IO.NodeOutput(audio)
|
||||
waveform = audio["waveform"]
|
||||
sample_rate = audio["sample_rate"]
|
||||
|
||||
gain = 10 ** (volume / 20)
|
||||
waveform = waveform * gain
|
||||
|
||||
return IO.NodeOutput({"waveform": waveform, "sample_rate": sample_rate})
|
||||
|
||||
adjust_volume = execute # TODO: remove
|
||||
|
||||
|
||||
class EmptyAudio(IO.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return IO.Schema(
|
||||
node_id="EmptyAudio",
|
||||
search_aliases=["blank audio"],
|
||||
display_name="Empty Audio",
|
||||
category="audio",
|
||||
inputs=[
|
||||
IO.Float.Input(
|
||||
"duration",
|
||||
default=60.0,
|
||||
min=0.0,
|
||||
max=0xffffffffffffffff,
|
||||
step=0.01,
|
||||
tooltip="Duration of the empty audio clip in seconds",
|
||||
),
|
||||
IO.Int.Input(
|
||||
"sample_rate",
|
||||
default=44100,
|
||||
tooltip="Sample rate of the empty audio clip.",
|
||||
min=1,
|
||||
max=192000,
|
||||
advanced=True,
|
||||
),
|
||||
IO.Int.Input(
|
||||
"channels",
|
||||
default=2,
|
||||
min=1,
|
||||
max=2,
|
||||
tooltip="Number of audio channels (1 for mono, 2 for stereo).",
|
||||
advanced=True,
|
||||
),
|
||||
],
|
||||
outputs=[IO.Audio.Output()],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, duration, sample_rate, channels) -> IO.NodeOutput:
|
||||
num_samples = int(round(duration * sample_rate))
|
||||
waveform = torch.zeros((1, channels, num_samples), dtype=torch.float32)
|
||||
return IO.NodeOutput({"waveform": waveform, "sample_rate": sample_rate})
|
||||
|
||||
create_empty_audio = execute # TODO: remove
|
||||
|
||||
|
||||
class AudioEqualizer3Band(IO.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return IO.Schema(
|
||||
node_id="AudioEqualizer3Band",
|
||||
search_aliases=["eq", "bass boost", "treble boost", "equalizer"],
|
||||
display_name="Audio Equalizer (3-Band)",
|
||||
category="audio",
|
||||
is_experimental=True,
|
||||
inputs=[
|
||||
IO.Audio.Input("audio"),
|
||||
IO.Float.Input("low_gain_dB", default=0.0, min=-24.0, max=24.0, step=0.1, tooltip="Gain for Low frequencies (Bass)"),
|
||||
IO.Int.Input("low_freq", default=100, min=20, max=500, tooltip="Cutoff frequency for Low shelf"),
|
||||
IO.Float.Input("mid_gain_dB", default=0.0, min=-24.0, max=24.0, step=0.1, tooltip="Gain for Mid frequencies"),
|
||||
IO.Int.Input("mid_freq", default=1000, min=200, max=4000, tooltip="Center frequency for Mids"),
|
||||
IO.Float.Input("mid_q", default=0.707, min=0.1, max=10.0, step=0.1, tooltip="Q factor (bandwidth) for Mids"),
|
||||
IO.Float.Input("high_gain_dB", default=0.0, min=-24.0, max=24.0, step=0.1, tooltip="Gain for High frequencies (Treble)"),
|
||||
IO.Int.Input("high_freq", default=5000, min=1000, max=15000, tooltip="Cutoff frequency for High shelf"),
|
||||
],
|
||||
outputs=[IO.Audio.Output()],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, audio, low_gain_dB, low_freq, mid_gain_dB, mid_freq, mid_q, high_gain_dB, high_freq) -> IO.NodeOutput:
|
||||
if audio is None:
|
||||
return IO.NodeOutput(None)
|
||||
waveform = audio["waveform"]
|
||||
sample_rate = audio["sample_rate"]
|
||||
|
||||
if waveform.shape[-1] == 0:
|
||||
return IO.NodeOutput(audio)
|
||||
|
||||
eq_waveform = waveform.clone()
|
||||
|
||||
# 1. Apply Low Shelf (Bass)
|
||||
if low_gain_dB != 0:
|
||||
eq_waveform = torchaudio.functional.bass_biquad(
|
||||
eq_waveform,
|
||||
sample_rate,
|
||||
gain=low_gain_dB,
|
||||
central_freq=float(low_freq),
|
||||
Q=0.707
|
||||
)
|
||||
|
||||
# 2. Apply Peaking EQ (Mids)
|
||||
if mid_gain_dB != 0:
|
||||
eq_waveform = torchaudio.functional.equalizer_biquad(
|
||||
eq_waveform,
|
||||
sample_rate,
|
||||
center_freq=float(mid_freq),
|
||||
gain=mid_gain_dB,
|
||||
Q=mid_q
|
||||
)
|
||||
|
||||
# 3. Apply High Shelf (Treble)
|
||||
if high_gain_dB != 0:
|
||||
eq_waveform = torchaudio.functional.treble_biquad(
|
||||
eq_waveform,
|
||||
sample_rate,
|
||||
gain=high_gain_dB,
|
||||
central_freq=float(high_freq),
|
||||
Q=0.707
|
||||
)
|
||||
|
||||
return IO.NodeOutput({"waveform": eq_waveform, "sample_rate": sample_rate})
|
||||
|
||||
|
||||
class AudioExtension(ComfyExtension):
|
||||
@override
|
||||
async def get_node_list(self) -> list[type[IO.ComfyNode]]:
|
||||
return [
|
||||
EmptyLatentAudio,
|
||||
VAEEncodeAudio,
|
||||
VAEDecodeAudio,
|
||||
VAEDecodeAudioTiled,
|
||||
SaveAudio,
|
||||
SaveAudioMP3,
|
||||
SaveAudioOpus,
|
||||
LoadAudio,
|
||||
PreviewAudio,
|
||||
ConditioningStableAudio,
|
||||
RecordAudio,
|
||||
TrimAudioDuration,
|
||||
SplitAudioChannels,
|
||||
JoinAudioChannels,
|
||||
AudioConcat,
|
||||
AudioMerge,
|
||||
AudioAdjustVolume,
|
||||
EmptyAudio,
|
||||
AudioEqualizer3Band,
|
||||
]
|
||||
|
||||
async def comfy_entrypoint() -> AudioExtension:
|
||||
return AudioExtension()
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,37 +0,0 @@
|
||||
comfyui-frontend-package==1.43.18
|
||||
comfyui-workflow-templates==0.9.77
|
||||
comfyui-embedded-docs==0.5.0
|
||||
torch
|
||||
torchsde
|
||||
torchvision
|
||||
torchaudio
|
||||
numpy>=1.25.0
|
||||
einops
|
||||
transformers>=4.50.3
|
||||
tokenizers>=0.13.3
|
||||
sentencepiece
|
||||
safetensors>=0.4.2
|
||||
aiohttp>=3.11.8
|
||||
yarl>=1.18.0
|
||||
pyyaml
|
||||
Pillow
|
||||
scipy
|
||||
tqdm
|
||||
psutil
|
||||
alembic
|
||||
SQLAlchemy>=2.0.0
|
||||
filelock
|
||||
av>=14.2.0
|
||||
comfy-kitchen>=0.2.8
|
||||
comfy-aimdo==0.3.0
|
||||
requests
|
||||
simpleeval>=1.0.0
|
||||
blake3
|
||||
|
||||
#non essential dependencies:
|
||||
kornia>=0.7.1
|
||||
spandrel
|
||||
pydantic~=2.0
|
||||
pydantic-settings~=2.0
|
||||
PyOpenGL
|
||||
glfw
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,289 +0,0 @@
|
||||
import argparse
|
||||
import pytest
|
||||
from requests.exceptions import HTTPError
|
||||
from unittest.mock import patch, mock_open
|
||||
|
||||
from app.frontend_management import (
|
||||
FrontendManager,
|
||||
FrontEndProvider,
|
||||
Release,
|
||||
)
|
||||
from comfy.cli_args import DEFAULT_VERSION_STRING
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_releases():
|
||||
return [
|
||||
Release(
|
||||
id=1,
|
||||
tag_name="1.0.0",
|
||||
name="Release 1.0.0",
|
||||
prerelease=False,
|
||||
created_at="2022-01-01T00:00:00Z",
|
||||
published_at="2022-01-01T00:00:00Z",
|
||||
body="Release notes for 1.0.0",
|
||||
assets=[{"name": "dist.zip", "url": "https://example.com/dist.zip"}],
|
||||
),
|
||||
Release(
|
||||
id=2,
|
||||
tag_name="2.0.0",
|
||||
name="Release 2.0.0",
|
||||
prerelease=False,
|
||||
created_at="2022-02-01T00:00:00Z",
|
||||
published_at="2022-02-01T00:00:00Z",
|
||||
body="Release notes for 2.0.0",
|
||||
assets=[{"name": "dist.zip", "url": "https://example.com/dist.zip"}],
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_provider(mock_releases):
|
||||
provider = FrontEndProvider(
|
||||
owner="test-owner",
|
||||
repo="test-repo",
|
||||
)
|
||||
provider.all_releases = mock_releases
|
||||
provider.latest_release = mock_releases[1]
|
||||
FrontendManager.PROVIDERS = [provider]
|
||||
return provider
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def clear_cache():
|
||||
import utils.install_util
|
||||
import app.frontend_management
|
||||
|
||||
utils.install_util.PACKAGE_VERSIONS = {}
|
||||
app.frontend_management.COMFY_PACKAGE_VERSIONS = []
|
||||
|
||||
|
||||
def test_get_release(mock_provider, mock_releases):
|
||||
version = "1.0.0"
|
||||
release = mock_provider.get_release(version)
|
||||
assert release == mock_releases[0]
|
||||
|
||||
|
||||
def test_get_release_latest(mock_provider, mock_releases):
|
||||
version = "latest"
|
||||
release = mock_provider.get_release(version)
|
||||
assert release == mock_releases[1]
|
||||
|
||||
|
||||
def test_get_release_invalid_version(mock_provider):
|
||||
version = "invalid"
|
||||
with pytest.raises(ValueError):
|
||||
mock_provider.get_release(version)
|
||||
|
||||
|
||||
def test_init_frontend_default():
|
||||
version_string = DEFAULT_VERSION_STRING
|
||||
frontend_path = FrontendManager.init_frontend(version_string)
|
||||
assert frontend_path == FrontendManager.default_frontend_path()
|
||||
|
||||
|
||||
def test_init_frontend_invalid_version():
|
||||
version_string = "test-owner/test-repo@1.100.99"
|
||||
with pytest.raises(HTTPError):
|
||||
FrontendManager.init_frontend_unsafe(version_string)
|
||||
|
||||
|
||||
def test_init_frontend_invalid_provider():
|
||||
version_string = "invalid/invalid@latest"
|
||||
with pytest.raises(HTTPError):
|
||||
FrontendManager.init_frontend_unsafe(version_string)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_os_functions():
|
||||
with (
|
||||
patch("app.frontend_management.os.makedirs") as mock_makedirs,
|
||||
patch("app.frontend_management.os.listdir") as mock_listdir,
|
||||
patch("app.frontend_management.os.rmdir") as mock_rmdir,
|
||||
):
|
||||
mock_listdir.return_value = [] # Simulate empty directory
|
||||
yield mock_makedirs, mock_listdir, mock_rmdir
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_download():
|
||||
with patch("app.frontend_management.download_release_asset_zip") as mock:
|
||||
mock.side_effect = Exception("Download failed") # Simulate download failure
|
||||
yield mock
|
||||
|
||||
|
||||
def test_finally_block(mock_os_functions, mock_download, mock_provider):
|
||||
# Arrange
|
||||
mock_makedirs, mock_listdir, mock_rmdir = mock_os_functions
|
||||
version_string = "test-owner/test-repo@1.0.0"
|
||||
|
||||
# Act & Assert
|
||||
with pytest.raises(Exception):
|
||||
FrontendManager.init_frontend_unsafe(version_string, mock_provider)
|
||||
|
||||
# Assert
|
||||
mock_makedirs.assert_called_once()
|
||||
mock_download.assert_called_once()
|
||||
mock_listdir.assert_called_once()
|
||||
mock_rmdir.assert_called_once()
|
||||
|
||||
|
||||
def test_parse_version_string():
|
||||
version_string = "owner/repo@1.0.0"
|
||||
repo_owner, repo_name, version = FrontendManager.parse_version_string(
|
||||
version_string
|
||||
)
|
||||
assert repo_owner == "owner"
|
||||
assert repo_name == "repo"
|
||||
assert version == "1.0.0"
|
||||
|
||||
|
||||
def test_parse_version_string_invalid():
|
||||
version_string = "invalid"
|
||||
with pytest.raises(argparse.ArgumentTypeError):
|
||||
FrontendManager.parse_version_string(version_string)
|
||||
|
||||
|
||||
def test_init_frontend_default_with_mocks():
|
||||
# Arrange
|
||||
version_string = DEFAULT_VERSION_STRING
|
||||
|
||||
# Act
|
||||
with (
|
||||
patch("app.frontend_management.check_comfy_packages_versions") as mock_check,
|
||||
patch.object(
|
||||
FrontendManager, "default_frontend_path", return_value="/mocked/path"
|
||||
),
|
||||
):
|
||||
frontend_path = FrontendManager.init_frontend(version_string)
|
||||
|
||||
# Assert
|
||||
assert frontend_path == "/mocked/path"
|
||||
mock_check.assert_called_once()
|
||||
|
||||
|
||||
def test_init_frontend_fallback_on_error():
|
||||
# Arrange
|
||||
version_string = "test-owner/test-repo@1.0.0"
|
||||
|
||||
# Act
|
||||
with (
|
||||
patch.object(
|
||||
FrontendManager, "init_frontend_unsafe", side_effect=Exception("Test error")
|
||||
),
|
||||
patch("app.frontend_management.check_comfy_packages_versions") as mock_check,
|
||||
patch.object(
|
||||
FrontendManager, "default_frontend_path", return_value="/default/path"
|
||||
),
|
||||
):
|
||||
frontend_path = FrontendManager.init_frontend(version_string)
|
||||
|
||||
# Assert
|
||||
assert frontend_path == "/default/path"
|
||||
mock_check.assert_called_once()
|
||||
|
||||
|
||||
def test_get_frontend_version():
|
||||
# Arrange
|
||||
expected_version = "1.25.0"
|
||||
mock_requirements_content = """torch
|
||||
torchsde
|
||||
comfyui-frontend-package==1.25.0
|
||||
other-package==1.0.0
|
||||
numpy"""
|
||||
|
||||
# Act
|
||||
with patch("builtins.open", mock_open(read_data=mock_requirements_content)):
|
||||
version = FrontendManager.get_required_frontend_version()
|
||||
|
||||
# Assert
|
||||
assert version == expected_version
|
||||
|
||||
|
||||
def test_get_frontend_version_invalid_semver():
|
||||
# Arrange
|
||||
mock_requirements_content = """torch
|
||||
torchsde
|
||||
comfyui-frontend-package==1.29.3.75
|
||||
other-package==1.0.0
|
||||
numpy"""
|
||||
|
||||
# Act
|
||||
with patch("builtins.open", mock_open(read_data=mock_requirements_content)):
|
||||
version = FrontendManager.get_required_frontend_version()
|
||||
|
||||
# Assert
|
||||
assert version is None
|
||||
|
||||
|
||||
def test_get_templates_version():
|
||||
# Arrange
|
||||
expected_version = "0.1.41"
|
||||
mock_requirements_content = """torch
|
||||
torchsde
|
||||
comfyui-frontend-package==1.25.0
|
||||
comfyui-workflow-templates==0.1.41
|
||||
other-package==1.0.0
|
||||
numpy"""
|
||||
|
||||
# Act
|
||||
with patch("builtins.open", mock_open(read_data=mock_requirements_content)):
|
||||
version = FrontendManager.get_required_templates_version()
|
||||
|
||||
# Assert
|
||||
assert version == expected_version
|
||||
|
||||
|
||||
def test_get_templates_version_not_found():
|
||||
# Arrange
|
||||
mock_requirements_content = """torch
|
||||
torchsde
|
||||
comfyui-frontend-package==1.25.0
|
||||
other-package==1.0.0
|
||||
numpy"""
|
||||
|
||||
# Act
|
||||
with patch("builtins.open", mock_open(read_data=mock_requirements_content)):
|
||||
version = FrontendManager.get_required_templates_version()
|
||||
|
||||
# Assert
|
||||
assert version is None
|
||||
|
||||
|
||||
def test_get_templates_version_invalid_semver():
|
||||
# Arrange
|
||||
mock_requirements_content = """torch
|
||||
torchsde
|
||||
comfyui-workflow-templates==1.0.0.beta
|
||||
other-package==1.0.0
|
||||
numpy"""
|
||||
|
||||
# Act
|
||||
with patch("builtins.open", mock_open(read_data=mock_requirements_content)):
|
||||
version = FrontendManager.get_required_templates_version()
|
||||
|
||||
# Assert
|
||||
assert version is None
|
||||
|
||||
|
||||
def test_get_installed_templates_version():
|
||||
# Arrange
|
||||
expected_version = "0.1.40"
|
||||
|
||||
# Act
|
||||
with patch("app.frontend_management.version", return_value=expected_version):
|
||||
version = FrontendManager.get_installed_templates_version()
|
||||
|
||||
# Assert
|
||||
assert version == expected_version
|
||||
|
||||
|
||||
def test_get_installed_templates_version_not_installed():
|
||||
# Act
|
||||
with patch(
|
||||
"app.frontend_management.version", side_effect=Exception("Package not found")
|
||||
):
|
||||
version = FrontendManager.get_installed_templates_version()
|
||||
|
||||
# Assert
|
||||
assert version is None
|
||||
@@ -1,144 +0,0 @@
|
||||
from collections import defaultdict
|
||||
|
||||
import torch
|
||||
|
||||
from comfy.model_detection import detect_unet_config, model_config_from_unet_config
|
||||
import comfy.supported_models
|
||||
|
||||
|
||||
def _freeze(value):
|
||||
"""Recursively convert a value to a hashable form so configs can be
|
||||
compared/used as dict keys or set members."""
|
||||
if isinstance(value, dict):
|
||||
return frozenset((k, _freeze(v)) for k, v in value.items())
|
||||
if isinstance(value, (list, tuple)):
|
||||
return tuple(_freeze(v) for v in value)
|
||||
if isinstance(value, set):
|
||||
return frozenset(_freeze(v) for v in value)
|
||||
return value
|
||||
|
||||
|
||||
def _make_longcat_comfyui_sd():
|
||||
"""Minimal ComfyUI-format state dict for pre-converted LongCat-Image weights."""
|
||||
sd = {}
|
||||
H = 32 # Reduce hidden state dimension to reduce memory usage
|
||||
C_IN = 16
|
||||
C_CTX = 3584
|
||||
|
||||
sd["img_in.weight"] = torch.empty(H, C_IN * 4)
|
||||
sd["img_in.bias"] = torch.empty(H)
|
||||
sd["txt_in.weight"] = torch.empty(H, C_CTX)
|
||||
sd["txt_in.bias"] = torch.empty(H)
|
||||
|
||||
sd["time_in.in_layer.weight"] = torch.empty(H, 256)
|
||||
sd["time_in.in_layer.bias"] = torch.empty(H)
|
||||
sd["time_in.out_layer.weight"] = torch.empty(H, H)
|
||||
sd["time_in.out_layer.bias"] = torch.empty(H)
|
||||
|
||||
sd["final_layer.adaLN_modulation.1.weight"] = torch.empty(2 * H, H)
|
||||
sd["final_layer.adaLN_modulation.1.bias"] = torch.empty(2 * H)
|
||||
sd["final_layer.linear.weight"] = torch.empty(C_IN * 4, H)
|
||||
sd["final_layer.linear.bias"] = torch.empty(C_IN * 4)
|
||||
|
||||
for i in range(19):
|
||||
sd[f"double_blocks.{i}.img_attn.norm.key_norm.weight"] = torch.empty(128)
|
||||
sd[f"double_blocks.{i}.img_attn.qkv.weight"] = torch.empty(3 * H, H)
|
||||
sd[f"double_blocks.{i}.img_mod.lin.weight"] = torch.empty(H, H)
|
||||
for i in range(38):
|
||||
sd[f"single_blocks.{i}.modulation.lin.weight"] = torch.empty(H, H)
|
||||
|
||||
return sd
|
||||
|
||||
|
||||
def _make_flux_schnell_comfyui_sd():
|
||||
"""Minimal ComfyUI-format state dict for standard Flux Schnell."""
|
||||
sd = {}
|
||||
H = 32 # Reduce hidden state dimension to reduce memory usage
|
||||
C_IN = 16
|
||||
|
||||
sd["img_in.weight"] = torch.empty(H, C_IN * 4)
|
||||
sd["img_in.bias"] = torch.empty(H)
|
||||
sd["txt_in.weight"] = torch.empty(H, 4096)
|
||||
sd["txt_in.bias"] = torch.empty(H)
|
||||
|
||||
sd["double_blocks.0.img_attn.norm.key_norm.weight"] = torch.empty(128)
|
||||
sd["double_blocks.0.img_attn.qkv.weight"] = torch.empty(3 * H, H)
|
||||
sd["double_blocks.0.img_mod.lin.weight"] = torch.empty(H, H)
|
||||
|
||||
for i in range(19):
|
||||
sd[f"double_blocks.{i}.img_attn.norm.key_norm.weight"] = torch.empty(128)
|
||||
for i in range(38):
|
||||
sd[f"single_blocks.{i}.modulation.lin.weight"] = torch.empty(H, H)
|
||||
|
||||
return sd
|
||||
|
||||
|
||||
class TestModelDetection:
|
||||
"""Verify that first-match model detection selects the correct model
|
||||
based on list ordering and unet_config specificity."""
|
||||
|
||||
def test_longcat_before_schnell_in_models_list(self):
|
||||
"""LongCatImage must appear before FluxSchnell in the models list."""
|
||||
models = comfy.supported_models.models
|
||||
longcat_idx = next(i for i, m in enumerate(models) if m.__name__ == "LongCatImage")
|
||||
schnell_idx = next(i for i, m in enumerate(models) if m.__name__ == "FluxSchnell")
|
||||
assert longcat_idx < schnell_idx, (
|
||||
f"LongCatImage (index {longcat_idx}) must come before "
|
||||
f"FluxSchnell (index {schnell_idx}) in the models list"
|
||||
)
|
||||
|
||||
def test_longcat_comfyui_detected_as_longcat(self):
|
||||
sd = _make_longcat_comfyui_sd()
|
||||
unet_config = detect_unet_config(sd, "")
|
||||
assert unet_config is not None
|
||||
assert unet_config["image_model"] == "flux"
|
||||
assert unet_config["context_in_dim"] == 3584
|
||||
assert unet_config["vec_in_dim"] is None
|
||||
assert unet_config["guidance_embed"] is False
|
||||
assert unet_config["txt_ids_dims"] == [1, 2]
|
||||
|
||||
model_config = model_config_from_unet_config(unet_config, sd)
|
||||
assert model_config is not None
|
||||
assert type(model_config).__name__ == "LongCatImage"
|
||||
|
||||
def test_longcat_comfyui_keys_pass_through_unchanged(self):
|
||||
"""Pre-converted weights should not be transformed by process_unet_state_dict."""
|
||||
sd = _make_longcat_comfyui_sd()
|
||||
unet_config = detect_unet_config(sd, "")
|
||||
model_config = model_config_from_unet_config(unet_config, sd)
|
||||
|
||||
processed = model_config.process_unet_state_dict(dict(sd))
|
||||
assert "img_in.weight" in processed
|
||||
assert "txt_in.weight" in processed
|
||||
assert "time_in.in_layer.weight" in processed
|
||||
assert "final_layer.linear.weight" in processed
|
||||
|
||||
def test_flux_schnell_comfyui_detected_as_flux_schnell(self):
|
||||
sd = _make_flux_schnell_comfyui_sd()
|
||||
unet_config = detect_unet_config(sd, "")
|
||||
assert unet_config is not None
|
||||
assert unet_config["image_model"] == "flux"
|
||||
assert unet_config["context_in_dim"] == 4096
|
||||
assert unet_config["txt_ids_dims"] == []
|
||||
|
||||
model_config = model_config_from_unet_config(unet_config, sd)
|
||||
assert model_config is not None
|
||||
assert type(model_config).__name__ == "FluxSchnell"
|
||||
|
||||
def test_unet_config_and_required_keys_combination_is_unique(self):
|
||||
"""Each model in the registry must have a unique combination of
|
||||
``unet_config`` and ``required_keys``. If two models share the same
|
||||
combination, ``BASE.matches`` cannot disambiguate between them and the
|
||||
first one in the list will always win."""
|
||||
models = comfy.supported_models.models
|
||||
groups = defaultdict(list)
|
||||
for model in models:
|
||||
key = (_freeze(model.unet_config), _freeze(model.required_keys))
|
||||
groups[key].append(model.__name__)
|
||||
|
||||
duplicates = {k: names for k, names in groups.items() if len(names) > 1}
|
||||
assert not duplicates, (
|
||||
"Found models sharing the same (unet_config, required_keys) "
|
||||
"combination, which makes detection ambiguous: "
|
||||
+ "; ".join(", ".join(names) for names in duplicates.values())
|
||||
)
|
||||
-1
Submodule comfy_source deleted from 9afc945331
@@ -8,6 +8,11 @@ cors_origins = ["*"]
|
||||
url = "sqlite:boosapi.db?mode=rwc"
|
||||
max_connections = 10
|
||||
ssl = false
|
||||
# PostgreSQL preparation example:
|
||||
# url = "postgres://boosapi:password@127.0.0.1:5432/boosapi"
|
||||
# Note: this build can create/migrate PostgreSQL schemas and copy SQLite data
|
||||
# with `boosapi migrate-sqlite-to-postgres`; full runtime query compatibility
|
||||
# is completed in a follow-up SQL dialect pass.
|
||||
|
||||
[auth]
|
||||
jwt_secret = "boosapi-dev-secret-change-in-production"
|
||||
|
||||
@@ -1,37 +0,0 @@
|
||||
FROM python:3.11-slim
|
||||
|
||||
WORKDIR /workspace
|
||||
|
||||
# Use USTC mirror for faster apt
|
||||
RUN sed -i 's|deb.debian.org|mirrors.ustc.edu.cn|g' /etc/apt/sources.list.d/debian.sources 2>/dev/null || \
|
||||
sed -i 's|deb.debian.org|mirrors.ustc.edu.cn|g' /etc/apt/sources.list
|
||||
|
||||
# Install system deps
|
||||
RUN apt-get update && apt-get install -y --no-install-recommends \
|
||||
git curl \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
# Copy ComfyUI from local source (avoid slow GitHub clone)
|
||||
COPY comfy_source /workspace/ComfyUI
|
||||
|
||||
WORKDIR /workspace/ComfyUI
|
||||
|
||||
# Install git (needed for some ComfyUI custom node managers)
|
||||
RUN git init
|
||||
|
||||
# Python venv + deps (CPU-only, generation goes through BoosAPI)
|
||||
RUN python3 -m venv venv && \
|
||||
. venv/bin/activate && \
|
||||
pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple && \
|
||||
pip install --upgrade pip && \
|
||||
pip config set global.extra-index-url https://download.pytorch.org/whl/cpu && \
|
||||
pip install torch torchvision torchaudio && \
|
||||
grep -v kornia_rs requirements.txt | pip install -r /dev/stdin && \
|
||||
pip install six && \
|
||||
rm -rf ~/.cache/pip
|
||||
|
||||
ENV PATH="/workspace/ComfyUI/venv/bin:${PATH}"
|
||||
|
||||
EXPOSE 8188
|
||||
|
||||
CMD ["python3", "main.py", "--listen", "0.0.0.0", "--port", "8188", "--cpu"]
|
||||
@@ -1,124 +0,0 @@
|
||||
# syntax=docker/dockerfile:1.4
|
||||
# Multi-tenant ComfyUI image — one container per BoosAPI user.
|
||||
#
|
||||
# Lineage: based on the same Python 3.11 / CPU-only profile the legacy
|
||||
# bosapi-comfyui:ready image used. Differences from the singleton image:
|
||||
# 1. Bakes `comfy_patches/overlay` into the image (proxy-base
|
||||
# default + auth-header fallback, server/runtime hardening, docs)
|
||||
# so containers do NOT need a host bind-mount of `comfy_source/`.
|
||||
# 2. Bakes the BoosAPI custom nodes from `comfy/boosapi_comfy_nodes.py`
|
||||
# under `/workspace/ComfyUI/custom_nodes/` so workflows referencing
|
||||
# `BoosAPITextGen` etc. resolve without any host mount.
|
||||
# 3. Reads `BOOSAPI_BUILTIN_KEY`, `BOOSAPI_BASE`, `BOOSAPI_KEY` at
|
||||
# runtime via the entrypoint shim and forwards them as
|
||||
# `--comfy-api-base` and the documented env vars.
|
||||
# 4. No host source mounts — all ComfyUI code lives inside the image
|
||||
# so the container can run with `--read-only` and tmpfs `/tmp`.
|
||||
#
|
||||
# Build context MUST be the repo root (`docker build -f docker/Dockerfile.comfyui.multi-tenant .`).
|
||||
#
|
||||
# Requires BuildKit (Docker 20.10+ has it on by default; build-and-run.sh
|
||||
# also exports DOCKER_BUILDKIT=1 defensively). The `--mount=type=cache`
|
||||
# instructions below are no-ops without BuildKit.
|
||||
|
||||
FROM python:3.11-slim
|
||||
|
||||
WORKDIR /workspace
|
||||
|
||||
# USTC mirror for faster apt
|
||||
RUN sed -i 's|deb.debian.org|mirrors.ustc.edu.cn|g' /etc/apt/sources.list.d/debian.sources 2>/dev/null || \
|
||||
sed -i 's|deb.debian.org|mirrors.ustc.edu.cn|g' /etc/apt/sources.list
|
||||
|
||||
RUN apt-get update && apt-get install -y --no-install-recommends \
|
||||
git curl tini \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
# uv: 5-10x faster than pip on resolution + install. Installed from PyPI
|
||||
# via the Tsinghua mirror that's already used by the rest of the build,
|
||||
# instead of curling the upstream GitHub release tarball —
|
||||
# `github.com/astral-sh/uv/releases/...` is unreliable from mainland
|
||||
# networks (TLS handshake completes, then body read truncates with
|
||||
# `SSL_read: unexpected eof`). The PyPI wheel bundles the same self-
|
||||
# contained binary and drops `uv` into the system Python's scripts dir
|
||||
# (already on PATH). Pinned for reproducibility.
|
||||
RUN pip install --no-cache-dir \
|
||||
--index-url https://pypi.tuna.tsinghua.edu.cn/simple \
|
||||
uv==0.5.31 \
|
||||
&& uv --version
|
||||
|
||||
# ── Layer-caching strategy ────────────────────────────────────────────
|
||||
# Copy ONLY requirements.txt first. The heavy install layers (torch +
|
||||
# ComfyUI deps, ~200 MB combined) then live in their own cache lines
|
||||
# that only invalidate when requirements.txt actually changes. Most
|
||||
# rebuilds change other comfy_source/ files (workflow scripts, custom
|
||||
# nodes, etc.), so this swap turns a 5-10 minute reinstall into a
|
||||
# zero-network layer-cache hit.
|
||||
#
|
||||
# The full `COPY comfy_source` happens after the install. It overlays
|
||||
# the rest of ComfyUI on top of the venv directory (which lives at
|
||||
# /workspace/ComfyUI/venv) without disturbing it — venv/ isn't present
|
||||
# in the source tree, and requirements.txt is just rewritten with
|
||||
# identical content.
|
||||
COPY comfy_source/requirements.txt /workspace/ComfyUI/requirements.txt
|
||||
WORKDIR /workspace/ComfyUI
|
||||
|
||||
RUN uv venv venv --python python3.11
|
||||
ENV PATH="/workspace/ComfyUI/venv/bin:${PATH}"
|
||||
|
||||
# ── PyTorch CPU wheels ────────────────────────────────────────────────
|
||||
# ~200 MB combined and the single biggest network draw of the build.
|
||||
# Two reliability levers:
|
||||
# 1. Tsinghua pytorch-wheels mirror replaces upstream
|
||||
# `download.pytorch.org/whl/cpu`. Upstream is slow + flaky from
|
||||
# mainland networks (TLS handshake completes, body read truncates
|
||||
# with `SSL_read: unexpected eof`). Tsinghua mirrors the same
|
||||
# wheels and is documented as a drop-in `--extra-index-url`.
|
||||
# 2. BuildKit cache mount on /root/.cache/uv preserves the downloaded
|
||||
# wheels across rebuilds. Even if this layer is invalidated (e.g.
|
||||
# after a torch version bump), uv reuses cached wheels and only
|
||||
# fetches what changed.
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
UV_INDEX_URL=https://pypi.tuna.tsinghua.edu.cn/simple \
|
||||
UV_EXTRA_INDEX_URL=https://mirrors.tuna.tsinghua.edu.cn/pytorch-wheels/cpu/ \
|
||||
uv pip install --python venv/bin/python torch torchvision torchaudio
|
||||
|
||||
# ── ComfyUI deps ──────────────────────────────────────────────────────
|
||||
# kornia_rs is a Rust-built wheel that compiles from source on install,
|
||||
# adding minutes to the build. The BoosAPI workflows we care about
|
||||
# don't use it, so strip it. Cache mount handles re-downloads when
|
||||
# requirements.txt changes (most lines stay the same → most wheels
|
||||
# reuse from cache).
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
grep -v kornia_rs requirements.txt > /tmp/requirements.txt && \
|
||||
UV_INDEX_URL=https://pypi.tuna.tsinghua.edu.cn/simple \
|
||||
uv pip install --python venv/bin/python -r /tmp/requirements.txt && \
|
||||
UV_INDEX_URL=https://pypi.tuna.tsinghua.edu.cn/simple \
|
||||
uv pip install --python venv/bin/python six && \
|
||||
rm /tmp/requirements.txt
|
||||
|
||||
# ── Full ComfyUI source ───────────────────────────────────────────────
|
||||
# This overlay COPY brings in everything else under comfy_source/. The
|
||||
# venv at /workspace/ComfyUI/venv survives because the source tree
|
||||
# doesn't contain a venv/ directory; requirements.txt is overwritten
|
||||
# with identical bytes (no-op).
|
||||
COPY comfy_source /workspace/ComfyUI
|
||||
RUN git init
|
||||
|
||||
# ── Bake the BoosAPI patches ──────────────────────────────────────────
|
||||
# Keep the image self-contained: overlay patches are tracked in the parent
|
||||
# repo, while comfy_source remains a vanilla ComfyUI submodule.
|
||||
COPY comfy_patches/overlay /workspace/ComfyUI
|
||||
|
||||
# BoosAPI custom nodes (BoosAPITextGen, image gen, etc.).
|
||||
RUN mkdir -p /workspace/ComfyUI/custom_nodes/boosapi_nodes
|
||||
COPY comfy/boosapi_comfy_nodes.py \
|
||||
/workspace/ComfyUI/custom_nodes/boosapi_nodes/__init__.py
|
||||
|
||||
# Entrypoint shim — forwards env vars to ComfyUI's CLI flags.
|
||||
COPY docker/comfyui-entrypoint.sh /usr/local/bin/comfyui-entrypoint
|
||||
RUN chmod +x /usr/local/bin/comfyui-entrypoint
|
||||
|
||||
EXPOSE 8188
|
||||
|
||||
# tini reaps any subprocesses ComfyUI spawns; the shim execs into Python.
|
||||
ENTRYPOINT ["/usr/bin/tini", "--", "/usr/local/bin/comfyui-entrypoint"]
|
||||
@@ -1,34 +0,0 @@
|
||||
#!/usr/bin/env bash
|
||||
# Multi-tenant ComfyUI entrypoint.
|
||||
#
|
||||
# Reads the per-user secrets the ComfyUIPool injects:
|
||||
# BOOSAPI_BUILTIN_KEY the BoosAPI downstream-key secret used by the
|
||||
# patched _helpers.py for built-in OpenAI nodes.
|
||||
# BOOSAPI_BASE the BoosAPI loopback base, e.g. http://host:4000/v1
|
||||
# BOOSAPI_KEY the per-user proxy key for boosapi_comfy_nodes.
|
||||
#
|
||||
# Then execs ComfyUI with --listen 0.0.0.0 --port 8188 and the
|
||||
# --comfy-api-base flag pointing at the BoosAPI agent shim so all
|
||||
# built-in OpenAI nodes route through us.
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
: "${BOOSAPI_BASE:?BOOSAPI_BASE is required (BoosAPI /v1 base URL)}"
|
||||
: "${BOOSAPI_BUILTIN_KEY:?BOOSAPI_BUILTIN_KEY is required (downstream-key secret)}"
|
||||
|
||||
# BOOSAPI_KEY is per-user; tolerate it being unset (then BoosAPITextGen
|
||||
# nodes will fall back to BOOSAPI_BUILTIN_KEY via the helper patch).
|
||||
export BOOSAPI_KEY="${BOOSAPI_KEY:-$BOOSAPI_BUILTIN_KEY}"
|
||||
|
||||
# Strip a trailing /v1 for the agent base (the helper appends it back).
|
||||
BOOSAPI_HOST="${BOOSAPI_BASE%/v1}"
|
||||
BOOSAPI_HOST="${BOOSAPI_HOST%/}"
|
||||
COMFY_API_BASE="${BOOSAPI_HOST}/api/comfyui-agent"
|
||||
|
||||
cd /workspace/ComfyUI
|
||||
|
||||
exec python main.py \
|
||||
--listen 0.0.0.0 \
|
||||
--port 8188 \
|
||||
--cpu \
|
||||
--comfy-api-base "${COMFY_API_BASE}"
|
||||
@@ -1,22 +0,0 @@
|
||||
# docker-compose.comfyui.yml
|
||||
# 开发环境：在 BoosAPI 旁边运行 ComfyUI
|
||||
# 使用: docker compose -f docker/docker-compose.yml -f docker/docker-compose.comfyui.yml up -d
|
||||
services:
|
||||
comfyui:
|
||||
build:
|
||||
context: ..
|
||||
dockerfile: docker/Dockerfile.comfyui
|
||||
ports:
|
||||
- "127.0.0.1:8188:8188"
|
||||
volumes:
|
||||
- comfyui-data:/workspace
|
||||
- ../boosapi_comfy_nodes.py:/workspace/ComfyUI/custom_nodes/boosapi_comfy_nodes.py
|
||||
environment:
|
||||
- BOOSAPI_BASE=http://host.docker.internal:4000/v1
|
||||
- BOOSAPI_KEY=${BOOSAPI_KEY:-sk-demo-user}
|
||||
restart: unless-stopped
|
||||
extra_hosts:
|
||||
- "host.docker.internal:host-gateway"
|
||||
|
||||
volumes:
|
||||
comfyui-data:
|
||||
@@ -0,0 +1,64 @@
|
||||
-- BoosAPI Workflow Studio core tables.
|
||||
|
||||
CREATE TABLE IF NOT EXISTS workflows (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
public_id TEXT NOT NULL UNIQUE,
|
||||
user_id INTEGER,
|
||||
name TEXT NOT NULL,
|
||||
description TEXT,
|
||||
graph TEXT NOT NULL DEFAULT '{}',
|
||||
is_template INTEGER NOT NULL DEFAULT 0,
|
||||
visibility TEXT NOT NULL DEFAULT 'private',
|
||||
created_at TEXT NOT NULL DEFAULT (datetime('now')),
|
||||
updated_at TEXT NOT NULL DEFAULT (datetime('now')),
|
||||
deleted_at TEXT,
|
||||
FOREIGN KEY(user_id) REFERENCES users(id) ON DELETE SET NULL
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_workflows_user_updated
|
||||
ON workflows(user_id, updated_at DESC);
|
||||
CREATE INDEX IF NOT EXISTS idx_workflows_visibility
|
||||
ON workflows(visibility);
|
||||
CREATE INDEX IF NOT EXISTS idx_workflows_deleted
|
||||
ON workflows(deleted_at);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS workflow_runs (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
public_id TEXT NOT NULL UNIQUE,
|
||||
workflow_id INTEGER NOT NULL,
|
||||
user_id INTEGER,
|
||||
status TEXT NOT NULL,
|
||||
input_overrides TEXT NOT NULL DEFAULT '{}',
|
||||
final_outputs TEXT,
|
||||
error_message TEXT,
|
||||
started_at TEXT,
|
||||
completed_at TEXT,
|
||||
created_at TEXT NOT NULL DEFAULT (datetime('now')),
|
||||
FOREIGN KEY(workflow_id) REFERENCES workflows(id) ON DELETE CASCADE,
|
||||
FOREIGN KEY(user_id) REFERENCES users(id) ON DELETE SET NULL
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_workflow_runs_workflow_created
|
||||
ON workflow_runs(workflow_id, created_at DESC);
|
||||
CREATE INDEX IF NOT EXISTS idx_workflow_runs_user_created
|
||||
ON workflow_runs(user_id, created_at DESC);
|
||||
CREATE INDEX IF NOT EXISTS idx_workflow_runs_status
|
||||
ON workflow_runs(status);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS workflow_node_runs (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
run_id INTEGER NOT NULL,
|
||||
node_id TEXT NOT NULL,
|
||||
node_type TEXT NOT NULL,
|
||||
status TEXT NOT NULL,
|
||||
inputs TEXT NOT NULL DEFAULT '{}',
|
||||
outputs TEXT,
|
||||
error_message TEXT,
|
||||
started_at TEXT,
|
||||
completed_at TEXT,
|
||||
created_at TEXT NOT NULL DEFAULT (datetime('now')),
|
||||
FOREIGN KEY(run_id) REFERENCES workflow_runs(id) ON DELETE CASCADE
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_workflow_node_runs_run
|
||||
ON workflow_node_runs(run_id, id);
|
||||
@@ -0,0 +1,868 @@
|
||||
-- BoosAPI initial schema
|
||||
-- Generated from the existing SQLite database (hub.db)
|
||||
-- This migration creates all tables, indexes, and constraints.
|
||||
|
||||
-- ============================================================
|
||||
-- SITES
|
||||
-- ============================================================
|
||||
CREATE TABLE IF NOT EXISTS sites (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
name TEXT NOT NULL,
|
||||
url TEXT NOT NULL,
|
||||
platform TEXT NOT NULL,
|
||||
status TEXT NOT NULL DEFAULT 'active',
|
||||
api_key TEXT,
|
||||
proxy_url TEXT,
|
||||
use_system_proxy BOOLEAN DEFAULT FALSE,
|
||||
custom_headers TEXT,
|
||||
external_checkin_url TEXT,
|
||||
global_weight DOUBLE PRECISION DEFAULT 1,
|
||||
post_refresh_probe_enabled BOOLEAN DEFAULT FALSE,
|
||||
post_refresh_probe_model TEXT DEFAULT '',
|
||||
post_refresh_probe_scope TEXT DEFAULT 'single',
|
||||
post_refresh_probe_latency_threshold_ms BIGINT DEFAULT 0,
|
||||
is_pinned BOOLEAN DEFAULT FALSE,
|
||||
sort_order BIGINT DEFAULT 0,
|
||||
created_at TEXT DEFAULT (CURRENT_TIMESTAMP::TEXT),
|
||||
updated_at TEXT DEFAULT (CURRENT_TIMESTAMP::TEXT)
|
||||
);
|
||||
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS sites_platform_url_unique
|
||||
ON sites (platform, url);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS sites_status_idx
|
||||
ON sites (status);
|
||||
|
||||
|
||||
-- ============================================================
|
||||
-- USERS
|
||||
-- ============================================================
|
||||
CREATE TABLE IF NOT EXISTS users (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
username TEXT NOT NULL,
|
||||
email TEXT NOT NULL,
|
||||
password_hash TEXT NOT NULL,
|
||||
role TEXT NOT NULL DEFAULT 'user',
|
||||
status TEXT NOT NULL DEFAULT 'active',
|
||||
created_at TEXT DEFAULT (CURRENT_TIMESTAMP::TEXT),
|
||||
updated_at TEXT DEFAULT (CURRENT_TIMESTAMP::TEXT)
|
||||
);
|
||||
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS users_email_unique
|
||||
ON users (email);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS users_status_idx
|
||||
ON users (status);
|
||||
|
||||
|
||||
-- ============================================================
|
||||
-- TOKEN ROUTES
|
||||
-- ============================================================
|
||||
CREATE TABLE IF NOT EXISTS token_routes (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
model_pattern TEXT NOT NULL,
|
||||
model_mapping TEXT,
|
||||
display_name TEXT,
|
||||
display_icon TEXT,
|
||||
decision_snapshot TEXT,
|
||||
decision_refreshed_at TEXT,
|
||||
routing_strategy TEXT DEFAULT 'weighted',
|
||||
route_mode TEXT DEFAULT 'pattern',
|
||||
enabled BOOLEAN DEFAULT TRUE,
|
||||
created_at TEXT DEFAULT (CURRENT_TIMESTAMP::TEXT),
|
||||
updated_at TEXT DEFAULT (CURRENT_TIMESTAMP::TEXT)
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS token_routes_model_pattern_idx
|
||||
ON token_routes (model_pattern);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS token_routes_enabled_idx
|
||||
ON token_routes (enabled);
|
||||
|
||||
|
||||
-- ============================================================
|
||||
-- ACCOUNTS
|
||||
-- ============================================================
|
||||
CREATE TABLE IF NOT EXISTS accounts (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
site_id BIGINT NOT NULL,
|
||||
username TEXT,
|
||||
access_token TEXT NOT NULL,
|
||||
api_token TEXT,
|
||||
balance DOUBLE PRECISION DEFAULT 0,
|
||||
balance_used DOUBLE PRECISION DEFAULT 0,
|
||||
quota DOUBLE PRECISION DEFAULT 0,
|
||||
unit_cost DOUBLE PRECISION,
|
||||
value_score DOUBLE PRECISION DEFAULT 0,
|
||||
status TEXT DEFAULT 'active',
|
||||
checkin_enabled BOOLEAN DEFAULT TRUE,
|
||||
last_checkin_at TEXT,
|
||||
last_balance_refresh TEXT,
|
||||
extra_config TEXT,
|
||||
is_pinned BOOLEAN DEFAULT FALSE,
|
||||
sort_order BIGINT DEFAULT 0,
|
||||
oauth_provider TEXT,
|
||||
oauth_account_key TEXT,
|
||||
oauth_project_id TEXT,
|
||||
created_at TEXT DEFAULT (CURRENT_TIMESTAMP::TEXT),
|
||||
updated_at TEXT DEFAULT (CURRENT_TIMESTAMP::TEXT),
|
||||
FOREIGN KEY (site_id) REFERENCES sites(id) ON DELETE CASCADE
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS accounts_site_id_idx
|
||||
ON accounts (site_id);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS accounts_status_idx
|
||||
ON accounts (status);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS accounts_site_status_idx
|
||||
ON accounts (site_id, status);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS accounts_oauth_provider_idx
|
||||
ON accounts (oauth_provider);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS accounts_oauth_identity_idx
|
||||
ON accounts (oauth_provider, oauth_account_key, oauth_project_id);
|
||||
|
||||
|
||||
-- ============================================================
|
||||
-- DOWNSTREAM API KEYS
|
||||
-- ============================================================
|
||||
CREATE TABLE IF NOT EXISTS downstream_api_keys (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
name TEXT NOT NULL,
|
||||
key TEXT NOT NULL,
|
||||
description TEXT,
|
||||
group_name TEXT,
|
||||
tags TEXT,
|
||||
enabled BOOLEAN DEFAULT TRUE,
|
||||
expires_at TEXT,
|
||||
max_cost DOUBLE PRECISION,
|
||||
used_cost DOUBLE PRECISION DEFAULT 0,
|
||||
max_requests BIGINT,
|
||||
used_requests BIGINT DEFAULT 0,
|
||||
supported_models TEXT,
|
||||
allowed_route_ids TEXT,
|
||||
site_weight_multipliers TEXT,
|
||||
excluded_site_ids TEXT,
|
||||
excluded_credential_refs TEXT,
|
||||
last_used_at TEXT,
|
||||
user_id BIGINT REFERENCES users(id),
|
||||
created_at TEXT DEFAULT (CURRENT_TIMESTAMP::TEXT),
|
||||
updated_at TEXT DEFAULT (CURRENT_TIMESTAMP::TEXT)
|
||||
);
|
||||
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS downstream_api_keys_key_unique
|
||||
ON downstream_api_keys (key);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS downstream_api_keys_name_idx
|
||||
ON downstream_api_keys (name);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS downstream_api_keys_enabled_idx
|
||||
ON downstream_api_keys (enabled);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS downstream_api_keys_expires_at_idx
|
||||
ON downstream_api_keys (expires_at);
|
||||
|
||||
|
||||
-- ============================================================
|
||||
-- SITE DISABLED MODELS
|
||||
-- ============================================================
|
||||
CREATE TABLE IF NOT EXISTS site_disabled_models (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
site_id BIGINT NOT NULL,
|
||||
model_name TEXT NOT NULL,
|
||||
created_at TEXT DEFAULT (CURRENT_TIMESTAMP::TEXT),
|
||||
FOREIGN KEY (site_id) REFERENCES sites(id) ON DELETE CASCADE
|
||||
);
|
||||
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS site_disabled_models_site_model_unique
|
||||
ON site_disabled_models (site_id, model_name);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS site_disabled_models_site_id_idx
|
||||
ON site_disabled_models (site_id);
|
||||
|
||||
|
||||
-- ============================================================
|
||||
-- SITE API ENDPOINTS
|
||||
-- ============================================================
|
||||
CREATE TABLE IF NOT EXISTS site_api_endpoints (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
site_id BIGINT NOT NULL,
|
||||
url TEXT NOT NULL,
|
||||
enabled BOOLEAN DEFAULT TRUE,
|
||||
sort_order BIGINT DEFAULT 0,
|
||||
cooldown_until TEXT,
|
||||
last_selected_at TEXT,
|
||||
last_failed_at TEXT,
|
||||
last_failure_reason TEXT,
|
||||
created_at TEXT DEFAULT (CURRENT_TIMESTAMP::TEXT),
|
||||
updated_at TEXT DEFAULT (CURRENT_TIMESTAMP::TEXT),
|
||||
FOREIGN KEY (site_id) REFERENCES sites(id) ON DELETE CASCADE
|
||||
);
|
||||
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS site_api_endpoints_site_url_unique
|
||||
ON site_api_endpoints (site_id, url);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS site_api_endpoints_site_enabled_sort_idx
|
||||
ON site_api_endpoints (site_id, enabled, sort_order);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS site_api_endpoints_site_cooldown_idx
|
||||
ON site_api_endpoints (site_id, cooldown_until);
|
||||
|
||||
|
||||
-- ============================================================
|
||||
-- ACCOUNT TOKENS
|
||||
-- ============================================================
|
||||
CREATE TABLE IF NOT EXISTS account_tokens (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
account_id BIGINT NOT NULL,
|
||||
name TEXT NOT NULL,
|
||||
token TEXT NOT NULL,
|
||||
source TEXT DEFAULT 'manual',
|
||||
token_group TEXT,
|
||||
value_status TEXT NOT NULL DEFAULT 'ready',
|
||||
enabled BOOLEAN DEFAULT TRUE,
|
||||
is_default BOOLEAN DEFAULT FALSE,
|
||||
created_at TEXT DEFAULT (CURRENT_TIMESTAMP::TEXT),
|
||||
updated_at TEXT DEFAULT (CURRENT_TIMESTAMP::TEXT),
|
||||
FOREIGN KEY (account_id) REFERENCES accounts(id) ON DELETE CASCADE
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS account_tokens_account_id_idx
|
||||
ON account_tokens (account_id);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS account_tokens_account_enabled_idx
|
||||
ON account_tokens (account_id, enabled);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS account_tokens_enabled_idx
|
||||
ON account_tokens (enabled);
|
||||
|
||||
|
||||
-- ============================================================
|
||||
-- CHECKIN LOGS
|
||||
-- ============================================================
|
||||
CREATE TABLE IF NOT EXISTS checkin_logs (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
account_id BIGINT NOT NULL,
|
||||
status TEXT NOT NULL,
|
||||
message TEXT,
|
||||
reward TEXT,
|
||||
created_at TEXT DEFAULT (CURRENT_TIMESTAMP::TEXT),
|
||||
FOREIGN KEY (account_id) REFERENCES accounts(id) ON DELETE CASCADE
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS checkin_logs_account_created_at_idx
|
||||
ON checkin_logs (account_id, created_at);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS checkin_logs_created_at_idx
|
||||
ON checkin_logs (created_at);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS checkin_logs_status_idx
|
||||
ON checkin_logs (status);
|
||||
|
||||
|
||||
-- ============================================================
|
||||
-- MODEL AVAILABILITY
|
||||
-- ============================================================
|
||||
CREATE TABLE IF NOT EXISTS model_availability (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
account_id BIGINT NOT NULL,
|
||||
model_name TEXT NOT NULL,
|
||||
available BIGINT,
|
||||
latency_ms BIGINT,
|
||||
checked_at TEXT DEFAULT (CURRENT_TIMESTAMP::TEXT),
|
||||
is_manual BOOLEAN DEFAULT FALSE,
|
||||
FOREIGN KEY (account_id) REFERENCES accounts(id) ON DELETE CASCADE
|
||||
);
|
||||
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS model_availability_account_model_unique
|
||||
ON model_availability (account_id, model_name);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS model_availability_account_available_idx
|
||||
ON model_availability (account_id, available);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS model_availability_model_name_idx
|
||||
ON model_availability (model_name);
|
||||
|
||||
|
||||
-- ============================================================
|
||||
-- TOKEN MODEL AVAILABILITY
|
||||
-- ============================================================
|
||||
CREATE TABLE IF NOT EXISTS token_model_availability (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
token_id BIGINT NOT NULL,
|
||||
model_name TEXT NOT NULL,
|
||||
available BIGINT,
|
||||
latency_ms BIGINT,
|
||||
checked_at TEXT DEFAULT (CURRENT_TIMESTAMP::TEXT),
|
||||
FOREIGN KEY (token_id) REFERENCES account_tokens(id) ON DELETE CASCADE
|
||||
);
|
||||
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS token_model_availability_token_model_unique
|
||||
ON token_model_availability (token_id, model_name);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS token_model_availability_token_available_idx
|
||||
ON token_model_availability (token_id, available);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS token_model_availability_model_name_idx
|
||||
ON token_model_availability (model_name);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS token_model_availability_available_idx
|
||||
ON token_model_availability (available);
|
||||
|
||||
|
||||
-- ============================================================
|
||||
-- ROUTE CHANNELS
|
||||
-- ============================================================
|
||||
CREATE TABLE IF NOT EXISTS route_channels (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
route_id BIGINT NOT NULL,
|
||||
account_id BIGINT NOT NULL,
|
||||
token_id BIGINT,
|
||||
oauth_route_unit_id BIGINT,
|
||||
source_model TEXT,
|
||||
priority BIGINT DEFAULT 0,
|
||||
weight BIGINT DEFAULT 10,
|
||||
enabled BOOLEAN DEFAULT TRUE,
|
||||
manual_override BOOLEAN DEFAULT FALSE,
|
||||
success_count BIGINT DEFAULT 0,
|
||||
fail_count BIGINT DEFAULT 0,
|
||||
total_latency_ms BIGINT DEFAULT 0,
|
||||
total_cost DOUBLE PRECISION DEFAULT 0,
|
||||
last_used_at TEXT,
|
||||
last_selected_at TEXT,
|
||||
last_fail_at TEXT,
|
||||
consecutive_fail_count BIGINT NOT NULL DEFAULT 0,
|
||||
cooldown_level BIGINT NOT NULL DEFAULT 0,
|
||||
cooldown_until TEXT,
|
||||
FOREIGN KEY (route_id) REFERENCES token_routes(id) ON DELETE CASCADE,
|
||||
FOREIGN KEY (account_id) REFERENCES accounts(id) ON DELETE CASCADE,
|
||||
FOREIGN KEY (token_id) REFERENCES account_tokens(id) ON DELETE SET NULL
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS route_channels_route_id_idx
|
||||
ON route_channels (route_id);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS route_channels_account_id_idx
|
||||
ON route_channels (account_id);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS route_channels_token_id_idx
|
||||
ON route_channels (token_id);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS route_channels_route_enabled_idx
|
||||
ON route_channels (route_id, enabled);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS route_channels_route_token_idx
|
||||
ON route_channels (route_id, token_id);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS route_channels_oauth_route_unit_id_idx
|
||||
ON route_channels (oauth_route_unit_id);
|
||||
|
||||
|
||||
-- ============================================================
|
||||
-- ROUTE GROUP SOURCES
|
||||
-- ============================================================
|
||||
CREATE TABLE IF NOT EXISTS route_group_sources (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
group_route_id BIGINT NOT NULL,
|
||||
source_route_id BIGINT NOT NULL,
|
||||
FOREIGN KEY (group_route_id) REFERENCES token_routes(id) ON DELETE CASCADE,
|
||||
FOREIGN KEY (source_route_id) REFERENCES token_routes(id) ON DELETE CASCADE
|
||||
);
|
||||
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS route_group_sources_group_source_unique
|
||||
ON route_group_sources (group_route_id, source_route_id);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS route_group_sources_source_route_id_idx
|
||||
ON route_group_sources (source_route_id);
|
||||
|
||||
|
||||
-- ============================================================
|
||||
-- PROXY LOGS
|
||||
-- ============================================================
|
||||
CREATE TABLE IF NOT EXISTS proxy_logs (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
route_id BIGINT,
|
||||
channel_id BIGINT,
|
||||
account_id BIGINT,
|
||||
downstream_api_key_id BIGINT,
|
||||
model_requested TEXT,
|
||||
model_actual TEXT,
|
||||
status TEXT,
|
||||
http_status BIGINT,
|
||||
latency_ms BIGINT,
|
||||
first_byte_latency_ms BIGINT,
|
||||
prompt_tokens BIGINT,
|
||||
completion_tokens BIGINT,
|
||||
total_tokens BIGINT,
|
||||
estimated_cost DOUBLE PRECISION,
|
||||
error_message TEXT,
|
||||
retry_count BIGINT DEFAULT 0,
|
||||
is_stream BIGINT,
|
||||
client_family TEXT,
|
||||
client_app_id TEXT,
|
||||
client_app_name TEXT,
|
||||
client_confidence TEXT,
|
||||
billing_details TEXT,
|
||||
created_at TEXT DEFAULT (CURRENT_TIMESTAMP::TEXT)
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS proxy_logs_created_at_idx
|
||||
ON proxy_logs (created_at);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS proxy_logs_account_created_at_idx
|
||||
ON proxy_logs (account_id, created_at);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS proxy_logs_status_created_at_idx
|
||||
ON proxy_logs (status, created_at);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS proxy_logs_model_actual_created_at_idx
|
||||
ON proxy_logs (model_actual, created_at);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS proxy_logs_downstream_api_key_created_at_idx
|
||||
ON proxy_logs (downstream_api_key_id, created_at);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS proxy_logs_client_app_id_created_at_idx
|
||||
ON proxy_logs (client_app_id, created_at);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS proxy_logs_client_family_created_at_idx
|
||||
ON proxy_logs (client_family, created_at);
|
||||
|
||||
|
||||
-- ============================================================
|
||||
-- PROXY VIDEO TASKS
|
||||
-- ============================================================
|
||||
CREATE TABLE IF NOT EXISTS proxy_video_tasks (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
public_id TEXT NOT NULL,
|
||||
upstream_video_id TEXT NOT NULL,
|
||||
site_url TEXT NOT NULL,
|
||||
token_value TEXT NOT NULL,
|
||||
requested_model TEXT,
|
||||
actual_model TEXT,
|
||||
channel_id BIGINT,
|
||||
account_id BIGINT,
|
||||
status_snapshot TEXT,
|
||||
upstream_response_meta TEXT,
|
||||
last_upstream_status BIGINT,
|
||||
last_polled_at TEXT,
|
||||
created_at TEXT DEFAULT (CURRENT_TIMESTAMP::TEXT),
|
||||
updated_at TEXT DEFAULT (CURRENT_TIMESTAMP::TEXT)
|
||||
);
|
||||
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS proxy_video_tasks_public_id_unique
|
||||
ON proxy_video_tasks (public_id);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS proxy_video_tasks_upstream_video_id_idx
|
||||
ON proxy_video_tasks (upstream_video_id);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS proxy_video_tasks_created_at_idx
|
||||
ON proxy_video_tasks (created_at);
|
||||
|
||||
|
||||
-- ============================================================
|
||||
-- PROXY FILES
|
||||
-- ============================================================
|
||||
CREATE TABLE IF NOT EXISTS proxy_files (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
public_id TEXT NOT NULL,
|
||||
owner_type TEXT NOT NULL,
|
||||
owner_id TEXT NOT NULL,
|
||||
filename TEXT NOT NULL,
|
||||
mime_type TEXT NOT NULL,
|
||||
purpose TEXT,
|
||||
byte_size BIGINT NOT NULL,
|
||||
sha256 TEXT NOT NULL,
|
||||
content_base64 TEXT NOT NULL,
|
||||
created_at TEXT DEFAULT (CURRENT_TIMESTAMP::TEXT),
|
||||
updated_at TEXT DEFAULT (CURRENT_TIMESTAMP::TEXT),
|
||||
deleted_at TEXT
|
||||
);
|
||||
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS proxy_files_public_id_unique
|
||||
ON proxy_files (public_id);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS proxy_files_owner_lookup_idx
|
||||
ON proxy_files (owner_type, owner_id, deleted_at);
|
||||
|
||||
|
||||
-- ============================================================
|
||||
-- PROXY DEBUG TRACES
|
||||
-- ============================================================
|
||||
CREATE TABLE IF NOT EXISTS proxy_debug_traces (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
downstream_path TEXT NOT NULL,
|
||||
client_kind TEXT,
|
||||
session_id TEXT,
|
||||
trace_hint TEXT,
|
||||
requested_model TEXT,
|
||||
downstream_api_key_id BIGINT,
|
||||
request_headers_json TEXT,
|
||||
request_body_json TEXT,
|
||||
sticky_session_key TEXT,
|
||||
sticky_hit_channel_id BIGINT,
|
||||
selected_channel_id BIGINT,
|
||||
selected_route_id BIGINT,
|
||||
selected_account_id BIGINT,
|
||||
selected_site_id BIGINT,
|
||||
selected_site_platform TEXT,
|
||||
endpoint_candidates_json TEXT,
|
||||
endpoint_runtime_state_json TEXT,
|
||||
decision_summary_json TEXT,
|
||||
final_status TEXT,
|
||||
final_http_status BIGINT,
|
||||
final_upstream_path TEXT,
|
||||
final_response_headers_json TEXT,
|
||||
final_response_body_json TEXT,
|
||||
created_at TEXT DEFAULT (CURRENT_TIMESTAMP::TEXT),
|
||||
updated_at TEXT DEFAULT (CURRENT_TIMESTAMP::TEXT)
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS proxy_debug_traces_created_at_idx
|
||||
ON proxy_debug_traces (created_at);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS proxy_debug_traces_session_created_at_idx
|
||||
ON proxy_debug_traces (session_id, created_at);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS proxy_debug_traces_model_created_at_idx
|
||||
ON proxy_debug_traces (requested_model, created_at);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS proxy_debug_traces_final_status_created_at_idx
|
||||
ON proxy_debug_traces (final_status, created_at);
|
||||
|
||||
|
||||
-- ============================================================
|
||||
-- PROXY DEBUG ATTEMPTS
|
||||
-- ============================================================
|
||||
CREATE TABLE IF NOT EXISTS proxy_debug_attempts (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
trace_id BIGINT NOT NULL,
|
||||
attempt_index BIGINT NOT NULL,
|
||||
endpoint TEXT NOT NULL,
|
||||
request_path TEXT NOT NULL,
|
||||
target_url TEXT NOT NULL,
|
||||
runtime_executor TEXT,
|
||||
request_headers_json TEXT,
|
||||
request_body_json TEXT,
|
||||
response_status BIGINT,
|
||||
response_headers_json TEXT,
|
||||
response_body_json TEXT,
|
||||
raw_error_text TEXT,
|
||||
recover_applied BOOLEAN DEFAULT FALSE,
|
||||
downgrade_decision BOOLEAN DEFAULT FALSE,
|
||||
downgrade_reason TEXT,
|
||||
memory_write_json TEXT,
|
||||
created_at TEXT DEFAULT (CURRENT_TIMESTAMP::TEXT),
|
||||
FOREIGN KEY (trace_id) REFERENCES proxy_debug_traces(id) ON DELETE CASCADE
|
||||
);
|
||||
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS proxy_debug_attempts_trace_attempt_unique
|
||||
ON proxy_debug_attempts (trace_id, attempt_index);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS proxy_debug_attempts_trace_created_at_idx
|
||||
ON proxy_debug_attempts (trace_id, created_at);
|
||||
|
||||
|
||||
-- ============================================================
|
||||
-- EVENTS
|
||||
-- ============================================================
|
||||
CREATE TABLE IF NOT EXISTS events (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
type TEXT NOT NULL,
|
||||
title TEXT NOT NULL,
|
||||
message TEXT,
|
||||
level TEXT NOT NULL DEFAULT 'info',
|
||||
"read" BOOLEAN DEFAULT FALSE,
|
||||
related_id BIGINT,
|
||||
related_type TEXT,
|
||||
created_at TEXT DEFAULT (CURRENT_TIMESTAMP::TEXT)
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS events_read_created_at_idx
|
||||
ON events ("read", created_at);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS events_type_created_at_idx
|
||||
ON events (type, created_at);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS events_created_at_idx
|
||||
ON events (created_at);
|
||||
|
||||
|
||||
-- ============================================================
|
||||
-- SITE ANNOUNCEMENTS
|
||||
-- ============================================================
|
||||
CREATE TABLE IF NOT EXISTS site_announcements (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
site_id BIGINT NOT NULL,
|
||||
platform TEXT NOT NULL,
|
||||
source_key TEXT NOT NULL,
|
||||
title TEXT NOT NULL,
|
||||
content TEXT NOT NULL,
|
||||
level TEXT NOT NULL DEFAULT 'info',
|
||||
source_url TEXT,
|
||||
starts_at TEXT,
|
||||
ends_at TEXT,
|
||||
upstream_created_at TEXT,
|
||||
upstream_updated_at TEXT,
|
||||
first_seen_at TEXT DEFAULT (CURRENT_TIMESTAMP::TEXT),
|
||||
last_seen_at TEXT DEFAULT (CURRENT_TIMESTAMP::TEXT),
|
||||
read_at TEXT,
|
||||
dismissed_at TEXT,
|
||||
raw_payload TEXT,
|
||||
FOREIGN KEY (site_id) REFERENCES sites(id) ON DELETE CASCADE
|
||||
);
|
||||
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS site_announcements_site_source_key_unique
|
||||
ON site_announcements (site_id, source_key);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS site_announcements_site_id_first_seen_at_idx
|
||||
ON site_announcements (site_id, first_seen_at);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS site_announcements_read_at_idx
|
||||
ON site_announcements (read_at);
|
||||
|
||||
|
||||
-- ============================================================
|
||||
-- OAUTH ROUTE UNITS
|
||||
-- ============================================================
|
||||
CREATE TABLE IF NOT EXISTS oauth_route_units (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
site_id BIGINT NOT NULL,
|
||||
provider TEXT NOT NULL,
|
||||
name TEXT NOT NULL,
|
||||
strategy TEXT NOT NULL DEFAULT 'round_robin',
|
||||
enabled BOOLEAN DEFAULT TRUE,
|
||||
created_at TEXT DEFAULT (CURRENT_TIMESTAMP::TEXT),
|
||||
updated_at TEXT DEFAULT (CURRENT_TIMESTAMP::TEXT),
|
||||
FOREIGN KEY (site_id) REFERENCES sites(id) ON DELETE CASCADE
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS oauth_route_units_site_provider_idx
|
||||
ON oauth_route_units (site_id, provider);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS oauth_route_units_enabled_idx
|
||||
ON oauth_route_units (enabled);
|
||||
|
||||
|
||||
-- ============================================================
|
||||
-- OAUTH ROUTE UNIT MEMBERS
|
||||
-- ============================================================
|
||||
CREATE TABLE IF NOT EXISTS oauth_route_unit_members (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
unit_id BIGINT NOT NULL,
|
||||
account_id BIGINT NOT NULL,
|
||||
sort_order BIGINT DEFAULT 0,
|
||||
success_count BIGINT DEFAULT 0,
|
||||
fail_count BIGINT DEFAULT 0,
|
||||
total_latency_ms BIGINT DEFAULT 0,
|
||||
total_cost DOUBLE PRECISION DEFAULT 0,
|
||||
last_used_at TEXT,
|
||||
last_selected_at TEXT,
|
||||
last_fail_at TEXT,
|
||||
consecutive_fail_count BIGINT NOT NULL DEFAULT 0,
|
||||
cooldown_level BIGINT NOT NULL DEFAULT 0,
|
||||
cooldown_until TEXT,
|
||||
created_at TEXT DEFAULT (CURRENT_TIMESTAMP::TEXT),
|
||||
updated_at TEXT DEFAULT (CURRENT_TIMESTAMP::TEXT),
|
||||
FOREIGN KEY (unit_id) REFERENCES oauth_route_units(id) ON DELETE CASCADE,
|
||||
FOREIGN KEY (account_id) REFERENCES accounts(id) ON DELETE CASCADE
|
||||
);
|
||||
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS oauth_route_unit_members_unit_account_unique
|
||||
ON oauth_route_unit_members (unit_id, account_id);
|
||||
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS oauth_route_unit_members_account_unique
|
||||
ON oauth_route_unit_members (account_id);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS oauth_route_unit_members_unit_sort_idx
|
||||
ON oauth_route_unit_members (unit_id, sort_order);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS oauth_route_unit_members_unit_cooldown_idx
|
||||
ON oauth_route_unit_members (unit_id, cooldown_until);
|
||||
|
||||
|
||||
-- ============================================================
|
||||
-- SETTINGS
|
||||
-- ============================================================
|
||||
CREATE TABLE IF NOT EXISTS settings (
|
||||
key TEXT PRIMARY KEY NOT NULL,
|
||||
value TEXT
|
||||
);
|
||||
|
||||
|
||||
-- ============================================================
|
||||
-- ADMIN SNAPSHOTS
|
||||
-- ============================================================
|
||||
CREATE TABLE IF NOT EXISTS admin_snapshots (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
namespace TEXT NOT NULL,
|
||||
snapshot_key TEXT NOT NULL,
|
||||
payload TEXT NOT NULL,
|
||||
generated_at TEXT NOT NULL,
|
||||
expires_at TEXT NOT NULL,
|
||||
stale_until TEXT NOT NULL,
|
||||
created_at TEXT DEFAULT (CURRENT_TIMESTAMP::TEXT),
|
||||
updated_at TEXT DEFAULT (CURRENT_TIMESTAMP::TEXT)
|
||||
);
|
||||
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS admin_snapshots_namespace_key_unique
|
||||
ON admin_snapshots (namespace, snapshot_key);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS admin_snapshots_expires_at_idx
|
||||
ON admin_snapshots (expires_at);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS admin_snapshots_stale_until_idx
|
||||
ON admin_snapshots (stale_until);
|
||||
|
||||
|
||||
-- ============================================================
|
||||
-- ANALYTICS PROJECTION CHECKPOINTS
|
||||
-- ============================================================
|
||||
CREATE TABLE IF NOT EXISTS analytics_projection_checkpoints (
|
||||
projector_key TEXT PRIMARY KEY NOT NULL,
|
||||
time_zone TEXT NOT NULL DEFAULT 'Local',
|
||||
last_proxy_log_id BIGINT NOT NULL DEFAULT 0,
|
||||
watermark_created_at TEXT,
|
||||
recompute_from_id BIGINT,
|
||||
recompute_requested_at TEXT,
|
||||
recompute_reason TEXT,
|
||||
recompute_started_at TEXT,
|
||||
recompute_completed_at TEXT,
|
||||
last_projected_at TEXT,
|
||||
last_successful_at TEXT,
|
||||
last_error TEXT,
|
||||
lease_owner TEXT,
|
||||
lease_token TEXT,
|
||||
lease_expires_at TEXT,
|
||||
created_at TEXT DEFAULT (CURRENT_TIMESTAMP::TEXT),
|
||||
updated_at TEXT DEFAULT (CURRENT_TIMESTAMP::TEXT)
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS analytics_projection_checkpoints_recompute_from_id_idx
|
||||
ON analytics_projection_checkpoints (recompute_from_id);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS analytics_projection_checkpoints_lease_expires_at_idx
|
||||
ON analytics_projection_checkpoints (lease_expires_at);
|
||||
|
||||
|
||||
-- ============================================================
|
||||
-- MODEL DAY USAGE
|
||||
-- ============================================================
|
||||
CREATE TABLE IF NOT EXISTS model_day_usage (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
local_day TEXT NOT NULL,
|
||||
site_id BIGINT NOT NULL,
|
||||
model TEXT NOT NULL,
|
||||
total_calls BIGINT NOT NULL DEFAULT 0,
|
||||
success_calls BIGINT NOT NULL DEFAULT 0,
|
||||
failed_calls BIGINT NOT NULL DEFAULT 0,
|
||||
total_tokens BIGINT NOT NULL DEFAULT 0,
|
||||
total_spend DOUBLE PRECISION NOT NULL DEFAULT 0,
|
||||
total_latency_ms BIGINT NOT NULL DEFAULT 0,
|
||||
latency_count BIGINT NOT NULL DEFAULT 0,
|
||||
created_at TEXT DEFAULT (CURRENT_TIMESTAMP::TEXT),
|
||||
updated_at TEXT DEFAULT (CURRENT_TIMESTAMP::TEXT),
|
||||
FOREIGN KEY (site_id) REFERENCES sites(id) ON DELETE CASCADE,
|
||||
CONSTRAINT model_day_usage_non_negative CHECK (
|
||||
total_calls >= 0
|
||||
AND success_calls >= 0
|
||||
AND failed_calls >= 0
|
||||
AND total_tokens >= 0
|
||||
AND total_spend >= 0
|
||||
AND total_latency_ms >= 0
|
||||
AND latency_count >= 0
|
||||
)
|
||||
);
|
||||
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS model_day_usage_day_site_model_unique
|
||||
ON model_day_usage (local_day, site_id, model);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS model_day_usage_day_idx
|
||||
ON model_day_usage (local_day);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS model_day_usage_site_id_idx
|
||||
ON model_day_usage (site_id);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS model_day_usage_model_idx
|
||||
ON model_day_usage (model);
|
||||
|
||||
|
||||
-- ============================================================
|
||||
-- SITE DAY USAGE
|
||||
-- ============================================================
|
||||
CREATE TABLE IF NOT EXISTS site_day_usage (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
local_day TEXT NOT NULL,
|
||||
site_id BIGINT NOT NULL,
|
||||
total_calls BIGINT NOT NULL DEFAULT 0,
|
||||
success_calls BIGINT NOT NULL DEFAULT 0,
|
||||
failed_calls BIGINT NOT NULL DEFAULT 0,
|
||||
total_tokens BIGINT NOT NULL DEFAULT 0,
|
||||
total_summary_spend DOUBLE PRECISION NOT NULL DEFAULT 0,
|
||||
total_site_spend DOUBLE PRECISION NOT NULL DEFAULT 0,
|
||||
total_latency_ms BIGINT NOT NULL DEFAULT 0,
|
||||
latency_count BIGINT NOT NULL DEFAULT 0,
|
||||
created_at TEXT DEFAULT (CURRENT_TIMESTAMP::TEXT),
|
||||
updated_at TEXT DEFAULT (CURRENT_TIMESTAMP::TEXT),
|
||||
FOREIGN KEY (site_id) REFERENCES sites(id) ON DELETE CASCADE,
|
||||
CONSTRAINT site_day_usage_non_negative CHECK (
|
||||
total_calls >= 0
|
||||
AND success_calls >= 0
|
||||
AND failed_calls >= 0
|
||||
AND total_tokens >= 0
|
||||
AND total_summary_spend >= 0
|
||||
AND total_site_spend >= 0
|
||||
AND total_latency_ms >= 0
|
||||
AND latency_count >= 0
|
||||
)
|
||||
);
|
||||
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS site_day_usage_day_site_unique
|
||||
ON site_day_usage (local_day, site_id);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS site_day_usage_day_idx
|
||||
ON site_day_usage (local_day);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS site_day_usage_site_id_idx
|
||||
ON site_day_usage (site_id);
|
||||
|
||||
|
||||
-- ============================================================
|
||||
-- SITE HOUR USAGE
|
||||
-- ============================================================
|
||||
CREATE TABLE IF NOT EXISTS site_hour_usage (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
bucket_start_utc TEXT NOT NULL,
|
||||
site_id BIGINT NOT NULL,
|
||||
total_calls BIGINT NOT NULL DEFAULT 0,
|
||||
success_calls BIGINT NOT NULL DEFAULT 0,
|
||||
failed_calls BIGINT NOT NULL DEFAULT 0,
|
||||
total_tokens BIGINT NOT NULL DEFAULT 0,
|
||||
total_summary_spend DOUBLE PRECISION NOT NULL DEFAULT 0,
|
||||
total_site_spend DOUBLE PRECISION NOT NULL DEFAULT 0,
|
||||
total_latency_ms BIGINT NOT NULL DEFAULT 0,
|
||||
latency_count BIGINT NOT NULL DEFAULT 0,
|
||||
created_at TEXT DEFAULT (CURRENT_TIMESTAMP::TEXT),
|
||||
updated_at TEXT DEFAULT (CURRENT_TIMESTAMP::TEXT),
|
||||
FOREIGN KEY (site_id) REFERENCES sites(id) ON DELETE CASCADE,
|
||||
CONSTRAINT site_hour_usage_non_negative CHECK (
|
||||
total_calls >= 0
|
||||
AND success_calls >= 0
|
||||
AND failed_calls >= 0
|
||||
AND total_tokens >= 0
|
||||
AND total_summary_spend >= 0
|
||||
AND total_site_spend >= 0
|
||||
AND total_latency_ms >= 0
|
||||
AND latency_count >= 0
|
||||
)
|
||||
);
|
||||
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS site_hour_usage_hour_site_unique
|
||||
ON site_hour_usage (bucket_start_utc, site_id);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS site_hour_usage_hour_idx
|
||||
ON site_hour_usage (bucket_start_utc);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS site_hour_usage_site_id_idx
|
||||
ON site_hour_usage (site_id);
|
||||
@@ -0,0 +1,24 @@
|
||||
-- Background tasks tracking
|
||||
CREATE TABLE IF NOT EXISTS background_tasks (
|
||||
id TEXT PRIMARY KEY NOT NULL,
|
||||
task_type TEXT NOT NULL,
|
||||
title TEXT NOT NULL,
|
||||
status TEXT NOT NULL DEFAULT 'pending',
|
||||
message TEXT DEFAULT '',
|
||||
error TEXT,
|
||||
dedupe_key TEXT,
|
||||
success_title TEXT,
|
||||
failure_title TEXT,
|
||||
notify_on_success BOOLEAN DEFAULT TRUE,
|
||||
notify_on_failure BOOLEAN DEFAULT TRUE,
|
||||
result_json TEXT,
|
||||
progress BIGINT DEFAULT 0,
|
||||
progress_total BIGINT DEFAULT 0,
|
||||
created_at TEXT DEFAULT (CURRENT_TIMESTAMP::TEXT),
|
||||
updated_at TEXT DEFAULT (CURRENT_TIMESTAMP::TEXT),
|
||||
started_at TEXT,
|
||||
finished_at TEXT
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS background_tasks_status_idx
|
||||
ON background_tasks (status);
|
||||
@@ -0,0 +1,53 @@
|
||||
-- Organizations, members, and invitations for multi-tenant support
|
||||
|
||||
CREATE TABLE IF NOT EXISTS organizations (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
name TEXT NOT NULL,
|
||||
slug TEXT NOT NULL UNIQUE,
|
||||
plan TEXT NOT NULL DEFAULT 'free',
|
||||
status TEXT NOT NULL DEFAULT 'active',
|
||||
created_at TEXT DEFAULT (CURRENT_TIMESTAMP::TEXT),
|
||||
updated_at TEXT DEFAULT (CURRENT_TIMESTAMP::TEXT)
|
||||
);
|
||||
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS org_slug_idx
|
||||
ON organizations (slug);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS org_plan_idx
|
||||
ON organizations (plan);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS organization_members (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
org_id BIGINT NOT NULL REFERENCES organizations(id) ON DELETE CASCADE,
|
||||
user_id BIGINT NOT NULL REFERENCES users(id) ON DELETE CASCADE,
|
||||
role TEXT NOT NULL DEFAULT 'member',
|
||||
status TEXT NOT NULL DEFAULT 'active',
|
||||
joined_at TEXT DEFAULT (CURRENT_TIMESTAMP::TEXT)
|
||||
);
|
||||
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS org_member_unique
|
||||
ON organization_members (org_id, user_id);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS org_members_org_id_idx
|
||||
ON organization_members (org_id);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS org_members_user_id_idx
|
||||
ON organization_members (user_id);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS invitations (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
org_id BIGINT NOT NULL REFERENCES organizations(id) ON DELETE CASCADE,
|
||||
email TEXT NOT NULL,
|
||||
token TEXT NOT NULL UNIQUE,
|
||||
role TEXT NOT NULL DEFAULT 'member',
|
||||
invited_by BIGINT REFERENCES users(id),
|
||||
status TEXT NOT NULL DEFAULT 'pending',
|
||||
expires_at TEXT NOT NULL,
|
||||
created_at TEXT DEFAULT (CURRENT_TIMESTAMP::TEXT)
|
||||
);
|
||||
|
||||
CREATE UNIQUE INDEX IF NOT EXISTS invitations_token_unique
|
||||
ON invitations (token);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS invitations_status_idx
|
||||
ON invitations (status);
|
||||
@@ -0,0 +1,37 @@
|
||||
-- Usage tracking tables migrated from the Node.js backend.
|
||||
|
||||
CREATE TABLE IF NOT EXISTS usage_logs (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
downstream_key_id BIGINT,
|
||||
model TEXT,
|
||||
prompt_tokens BIGINT DEFAULT 0,
|
||||
completion_tokens BIGINT DEFAULT 0,
|
||||
total_tokens BIGINT DEFAULT 0,
|
||||
cost DOUBLE PRECISION DEFAULT 0.0,
|
||||
created_at TEXT DEFAULT (CURRENT_TIMESTAMP::TEXT),
|
||||
FOREIGN KEY (downstream_key_id) REFERENCES downstream_api_keys(id) ON DELETE SET NULL
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS usage_logs_downstream_key_created_idx
|
||||
ON usage_logs (downstream_key_id, created_at);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS usage_logs_created_idx
|
||||
ON usage_logs (created_at);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS usage_daily_snapshots (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
local_day TEXT NOT NULL,
|
||||
downstream_key_id BIGINT,
|
||||
request_count BIGINT NOT NULL DEFAULT 0,
|
||||
prompt_tokens BIGINT NOT NULL DEFAULT 0,
|
||||
completion_tokens BIGINT NOT NULL DEFAULT 0,
|
||||
total_tokens BIGINT NOT NULL DEFAULT 0,
|
||||
total_cost DOUBLE PRECISION NOT NULL DEFAULT 0.0,
|
||||
created_at TEXT DEFAULT (CURRENT_TIMESTAMP::TEXT),
|
||||
updated_at TEXT DEFAULT (CURRENT_TIMESTAMP::TEXT),
|
||||
UNIQUE(),
|
||||
FOREIGN KEY (downstream_key_id) REFERENCES downstream_api_keys(id) ON DELETE SET NULL
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS usage_daily_snapshots_day_idx
|
||||
ON usage_daily_snapshots (local_day);
|
||||
@@ -0,0 +1,33 @@
|
||||
-- Admin-authored announcements for normal users.
|
||||
|
||||
CREATE TABLE IF NOT EXISTS user_announcements (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
title TEXT NOT NULL,
|
||||
content TEXT NOT NULL,
|
||||
level TEXT NOT NULL DEFAULT 'info',
|
||||
status TEXT NOT NULL DEFAULT 'published',
|
||||
audience TEXT NOT NULL DEFAULT 'all',
|
||||
starts_at TEXT,
|
||||
ends_at TEXT,
|
||||
created_at TEXT NOT NULL DEFAULT (CURRENT_TIMESTAMP::TEXT),
|
||||
updated_at TEXT NOT NULL DEFAULT (CURRENT_TIMESTAMP::TEXT),
|
||||
published_at TEXT
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS user_announcements_status_created_at_idx
|
||||
ON user_announcements (status, created_at);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS user_announcements_active_window_idx
|
||||
ON user_announcements (starts_at, ends_at);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS user_announcement_reads (
|
||||
announcement_id BIGINT NOT NULL,
|
||||
user_id BIGINT NOT NULL,
|
||||
read_at TEXT NOT NULL DEFAULT (CURRENT_TIMESTAMP::TEXT),
|
||||
PRIMARY KEY (announcement_id, user_id),
|
||||
FOREIGN KEY (announcement_id) REFERENCES user_announcements(id) ON DELETE CASCADE,
|
||||
FOREIGN KEY (user_id) REFERENCES users(id) ON DELETE CASCADE
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS user_announcement_reads_user_id_idx
|
||||
ON user_announcement_reads (user_id);
|
||||
@@ -0,0 +1,34 @@
|
||||
-- Workspace assets: local archive for creative outputs from Picture Agent,
|
||||
-- ComfyUI Agent, Video Agent and ComfyUI workspaces.
|
||||
|
||||
CREATE TABLE IF NOT EXISTS workspace_assets (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
public_id TEXT NOT NULL UNIQUE,
|
||||
user_id BIGINT,
|
||||
source TEXT NOT NULL,
|
||||
kind TEXT NOT NULL,
|
||||
title TEXT,
|
||||
prompt TEXT,
|
||||
model TEXT,
|
||||
status TEXT NOT NULL DEFAULT 'completed',
|
||||
mime_type TEXT,
|
||||
file_name TEXT,
|
||||
file_path TEXT,
|
||||
remote_url TEXT,
|
||||
thumbnail_path TEXT,
|
||||
size_bytes BIGINT NOT NULL DEFAULT 0,
|
||||
metadata TEXT NOT NULL DEFAULT '{}',
|
||||
created_at TEXT NOT NULL DEFAULT (CURRENT_TIMESTAMP::TEXT),
|
||||
updated_at TEXT NOT NULL DEFAULT (CURRENT_TIMESTAMP::TEXT),
|
||||
deleted_at TEXT,
|
||||
FOREIGN KEY(user_id) REFERENCES users(id) ON DELETE SET NULL
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_workspace_assets_user_created
|
||||
ON workspace_assets(user_id, created_at DESC);
|
||||
CREATE INDEX IF NOT EXISTS idx_workspace_assets_source_created
|
||||
ON workspace_assets(source, created_at DESC);
|
||||
CREATE INDEX IF NOT EXISTS idx_workspace_assets_kind_created
|
||||
ON workspace_assets(kind, created_at DESC);
|
||||
CREATE INDEX IF NOT EXISTS idx_workspace_assets_deleted
|
||||
ON workspace_assets(deleted_at);
|
||||
@@ -0,0 +1,64 @@
|
||||
-- BoosAPI Workflow Studio core tables.
|
||||
|
||||
CREATE TABLE IF NOT EXISTS workflows (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
public_id TEXT NOT NULL UNIQUE,
|
||||
user_id BIGINT,
|
||||
name TEXT NOT NULL,
|
||||
description TEXT,
|
||||
graph TEXT NOT NULL DEFAULT '{}',
|
||||
is_template BIGINT NOT NULL DEFAULT 0,
|
||||
visibility TEXT NOT NULL DEFAULT 'private',
|
||||
created_at TEXT NOT NULL DEFAULT (CURRENT_TIMESTAMP::TEXT),
|
||||
updated_at TEXT NOT NULL DEFAULT (CURRENT_TIMESTAMP::TEXT),
|
||||
deleted_at TEXT,
|
||||
FOREIGN KEY(user_id) REFERENCES users(id) ON DELETE SET NULL
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_workflows_user_updated
|
||||
ON workflows(user_id, updated_at DESC);
|
||||
CREATE INDEX IF NOT EXISTS idx_workflows_visibility
|
||||
ON workflows(visibility);
|
||||
CREATE INDEX IF NOT EXISTS idx_workflows_deleted
|
||||
ON workflows(deleted_at);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS workflow_runs (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
public_id TEXT NOT NULL UNIQUE,
|
||||
workflow_id BIGINT NOT NULL,
|
||||
user_id BIGINT,
|
||||
status TEXT NOT NULL,
|
||||
input_overrides TEXT NOT NULL DEFAULT '{}',
|
||||
final_outputs TEXT,
|
||||
error_message TEXT,
|
||||
started_at TEXT,
|
||||
completed_at TEXT,
|
||||
created_at TEXT NOT NULL DEFAULT (CURRENT_TIMESTAMP::TEXT),
|
||||
FOREIGN KEY(workflow_id) REFERENCES workflows(id) ON DELETE CASCADE,
|
||||
FOREIGN KEY(user_id) REFERENCES users(id) ON DELETE SET NULL
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_workflow_runs_workflow_created
|
||||
ON workflow_runs(workflow_id, created_at DESC);
|
||||
CREATE INDEX IF NOT EXISTS idx_workflow_runs_user_created
|
||||
ON workflow_runs(user_id, created_at DESC);
|
||||
CREATE INDEX IF NOT EXISTS idx_workflow_runs_status
|
||||
ON workflow_runs(status);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS workflow_node_runs (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
run_id BIGINT NOT NULL,
|
||||
node_id TEXT NOT NULL,
|
||||
node_type TEXT NOT NULL,
|
||||
status TEXT NOT NULL,
|
||||
inputs TEXT NOT NULL DEFAULT '{}',
|
||||
outputs TEXT,
|
||||
error_message TEXT,
|
||||
started_at TEXT,
|
||||
completed_at TEXT,
|
||||
created_at TEXT NOT NULL DEFAULT (CURRENT_TIMESTAMP::TEXT),
|
||||
FOREIGN KEY(run_id) REFERENCES workflow_runs(id) ON DELETE CASCADE
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_workflow_node_runs_run
|
||||
ON workflow_node_runs(run_id, id);
|
||||
+2
-1
@@ -30,7 +30,7 @@
|
||||
"dev": "concurrently \"npm run dev:server\" \"vite\"",
|
||||
"dev:desktop": "concurrently \"npm run dev:server\" \"vite\" \"tsc -p tsconfig.desktop.json --watch --preserveWatchOutput\" \"wait-on tcp:127.0.0.1:4000 http://127.0.0.1:5173 file:dist/desktop/main.js && cross-env METAPI_DESKTOP_DEV_SERVER_URL=http://127.0.0.1:5173 METAPI_DESKTOP_EXTERNAL_SERVER_URL=http://127.0.0.1:4000 electron dist/desktop/main.js\"",
|
||||
"desktop:icons": "node scripts/desktop/generate-icons.mjs",
|
||||
"build:web": "npm run desktop:icons && vite build",
|
||||
"build:web": "vite build",
|
||||
"build:server": "tsc -p tsconfig.server.json && tsx scripts/dev/copy-runtime-db-generated.ts",
|
||||
"build:desktop": "tsc -p tsconfig.desktop.json",
|
||||
"build": "npm run build:web && npm run build:server && npm run build:desktop",
|
||||
@@ -52,6 +52,7 @@
|
||||
"comfyui:install": "node scripts/comfyui/install.mjs",
|
||||
"comfyui:start": "node scripts/comfyui/start.mjs",
|
||||
"comfyui:stop": "node scripts/comfyui/stop.mjs",
|
||||
"workflow-ui:install": "node scripts/workflow-ui/install.mjs",
|
||||
"db:generate": "drizzle-kit generate",
|
||||
"db:migrate": "tsx src/server/db/migrate.ts",
|
||||
"schema:contract": "tsx scripts/dev/generate-schema-contract.ts",
|
||||
|
||||
@@ -0,0 +1,34 @@
|
||||
import fs from 'fs/promises';
|
||||
import path from 'path';
|
||||
import { spawn } from 'child_process';
|
||||
|
||||
const ROOT = process.cwd();
|
||||
const REPO = process.env.BOOSAPI_WORKFLOW_UI_REPO || 'https://github.com/Comfy-Org/ComfyUI_frontend.git';
|
||||
const REF = process.env.BOOSAPI_WORKFLOW_UI_REF || 'v1.9.19';
|
||||
const DIR = path.join(ROOT, '.cache', 'comfyui-frontend');
|
||||
const DIST = path.join(ROOT, 'dist', 'workflow-ui');
|
||||
|
||||
function run(command, args, options = {}) {
|
||||
return new Promise((resolve, reject) => {
|
||||
const child = spawn(command, args, {
|
||||
cwd: options.cwd || ROOT,
|
||||
env: { ...process.env, ...(options.env || {}) },
|
||||
stdio: 'inherit',
|
||||
shell: process.platform === 'win32',
|
||||
});
|
||||
child.on('error', reject);
|
||||
child.on('exit', (code) => {
|
||||
if (code === 0) resolve();
|
||||
else reject(new Error(`${command} ${args.join(' ')} exited with ${code}`));
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
await fs.mkdir(path.dirname(DIR), { recursive: true });
|
||||
await fs.rm(DIR, { recursive: true, force: true });
|
||||
await run('git', ['clone', '--depth', '1', '--branch', REF, REPO, DIR]);
|
||||
await run('npm', ['ci'], { cwd: DIR });
|
||||
await run('npm', ['run', 'build'], { cwd: DIR });
|
||||
await fs.rm(DIST, { recursive: true, force: true });
|
||||
await fs.cp(path.join(DIR, 'dist'), DIST, { recursive: true });
|
||||
console.log(`workflow-ui installed to ${DIST}`);
|
||||
+100
-1
@@ -1,12 +1,111 @@
|
||||
use sqlx::sqlite::{SqlitePool, SqlitePoolOptions};
|
||||
use anyhow::{bail, Context};
|
||||
use serde::Serialize;
|
||||
use sqlx::{
|
||||
any::{install_default_drivers, AnyPoolOptions},
|
||||
postgres::{PgPool, PgPoolOptions},
|
||||
sqlite::{SqlitePool, SqlitePoolOptions},
|
||||
};
|
||||
|
||||
pub mod models;
|
||||
pub mod queries;
|
||||
pub mod sqlite_to_postgres;
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum DbKind {
|
||||
Sqlite,
|
||||
Postgres,
|
||||
}
|
||||
|
||||
impl DbKind {
|
||||
pub fn from_url(database_url: &str) -> anyhow::Result<Self> {
|
||||
let normalized = database_url.trim().to_ascii_lowercase();
|
||||
if normalized.starts_with("sqlite:") {
|
||||
Ok(Self::Sqlite)
|
||||
} else if normalized.starts_with("postgres://") || normalized.starts_with("postgresql://") {
|
||||
Ok(Self::Postgres)
|
||||
} else {
|
||||
bail!("unsupported database URL; expected sqlite:, postgres:// or postgresql://")
|
||||
}
|
||||
}
|
||||
|
||||
pub fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
Self::Sqlite => "sqlite",
|
||||
Self::Postgres => "postgres",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn create_pool(database_url: &str, max_connections: u32) -> anyhow::Result<SqlitePool> {
|
||||
if DbKind::from_url(database_url)? != DbKind::Sqlite {
|
||||
bail!(
|
||||
"runtime PostgreSQL query support is not enabled yet; use migrate-sqlite-to-postgres to prepare PostgreSQL data"
|
||||
);
|
||||
}
|
||||
let pool = SqlitePoolOptions::new()
|
||||
.max_connections(max_connections)
|
||||
.connect(database_url)
|
||||
.await?;
|
||||
Ok(pool)
|
||||
}
|
||||
|
||||
pub async fn create_pg_pool(database_url: &str, max_connections: u32) -> anyhow::Result<PgPool> {
|
||||
let pool = PgPoolOptions::new()
|
||||
.max_connections(max_connections)
|
||||
.connect(database_url)
|
||||
.await
|
||||
.context("connect PostgreSQL")?;
|
||||
Ok(pool)
|
||||
}
|
||||
|
||||
pub async fn test_connection(database_url: &str, max_connections: u32) -> anyhow::Result<DbKind> {
|
||||
match DbKind::from_url(database_url)? {
|
||||
DbKind::Sqlite => {
|
||||
SqlitePoolOptions::new()
|
||||
.max_connections(max_connections)
|
||||
.connect(database_url)
|
||||
.await
|
||||
.context("connect SQLite")?
|
||||
.close()
|
||||
.await;
|
||||
Ok(DbKind::Sqlite)
|
||||
}
|
||||
DbKind::Postgres => {
|
||||
create_pg_pool(database_url, max_connections)
|
||||
.await?
|
||||
.close()
|
||||
.await;
|
||||
Ok(DbKind::Postgres)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn run_migrations(
|
||||
kind: DbKind,
|
||||
sqlite_pool: Option<&SqlitePool>,
|
||||
pg_pool: Option<&PgPool>,
|
||||
) -> anyhow::Result<()> {
|
||||
match kind {
|
||||
DbKind::Sqlite => {
|
||||
let pool = sqlite_pool.context("SQLite pool required for SQLite migrations")?;
|
||||
sqlx::migrate!("./migrations").run(pool).await?;
|
||||
}
|
||||
DbKind::Postgres => {
|
||||
let pool = pg_pool.context("PostgreSQL pool required for PostgreSQL migrations")?;
|
||||
sqlx::migrate!("./migrations_pg").run(pool).await?;
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn create_any_pool(
|
||||
database_url: &str,
|
||||
max_connections: u32,
|
||||
) -> anyhow::Result<sqlx::AnyPool> {
|
||||
install_default_drivers();
|
||||
Ok(AnyPoolOptions::new()
|
||||
.max_connections(max_connections)
|
||||
.connect(database_url)
|
||||
.await?)
|
||||
}
|
||||
|
||||
@@ -0,0 +1,463 @@
|
||||
use anyhow::{bail, Context};
|
||||
use serde::Serialize;
|
||||
use sqlx::{
|
||||
postgres::PgPool,
|
||||
sqlite::{SqlitePool, SqliteRow},
|
||||
Row,
|
||||
};
|
||||
use std::collections::HashMap;
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct MigrationSummary {
|
||||
pub tables: Vec<TableMigrationSummary>,
|
||||
pub total_rows: i64,
|
||||
pub force_truncate: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct TableMigrationSummary {
|
||||
pub table: &'static str,
|
||||
pub rows: i64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
enum ColumnKind {
|
||||
I64,
|
||||
F64,
|
||||
Bool,
|
||||
Text,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
struct ColumnSpec {
|
||||
name: &'static str,
|
||||
kind: ColumnKind,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
struct TableSpec {
|
||||
name: &'static str,
|
||||
columns: &'static [ColumnSpec],
|
||||
}
|
||||
|
||||
macro_rules! col {
|
||||
($name:literal, i64) => {
|
||||
ColumnSpec {
|
||||
name: $name,
|
||||
kind: ColumnKind::I64,
|
||||
}
|
||||
};
|
||||
($name:literal, f64) => {
|
||||
ColumnSpec {
|
||||
name: $name,
|
||||
kind: ColumnKind::F64,
|
||||
}
|
||||
};
|
||||
($name:literal, bool) => {
|
||||
ColumnSpec {
|
||||
name: $name,
|
||||
kind: ColumnKind::Bool,
|
||||
}
|
||||
};
|
||||
($name:literal, text) => {
|
||||
ColumnSpec {
|
||||
name: $name,
|
||||
kind: ColumnKind::Text,
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
const TABLES: &[TableSpec] = &[
|
||||
TableSpec {
|
||||
name: "users",
|
||||
columns: &[
|
||||
col!("id", i64),
|
||||
col!("username", text),
|
||||
col!("email", text),
|
||||
col!("password_hash", text),
|
||||
col!("role", text),
|
||||
col!("status", text),
|
||||
col!("created_at", text),
|
||||
col!("updated_at", text),
|
||||
],
|
||||
},
|
||||
TableSpec {
|
||||
name: "sites",
|
||||
columns: &[
|
||||
col!("id", i64),
|
||||
col!("name", text),
|
||||
col!("url", text),
|
||||
col!("platform", text),
|
||||
col!("status", text),
|
||||
col!("api_key", text),
|
||||
col!("proxy_url", text),
|
||||
col!("use_system_proxy", bool),
|
||||
col!("custom_headers", text),
|
||||
col!("external_checkin_url", text),
|
||||
col!("global_weight", f64),
|
||||
col!("post_refresh_probe_enabled", bool),
|
||||
col!("post_refresh_probe_model", text),
|
||||
col!("post_refresh_probe_scope", text),
|
||||
col!("post_refresh_probe_latency_threshold_ms", i64),
|
||||
col!("is_pinned", bool),
|
||||
col!("sort_order", i64),
|
||||
col!("created_at", text),
|
||||
col!("updated_at", text),
|
||||
],
|
||||
},
|
||||
TableSpec {
|
||||
name: "accounts",
|
||||
columns: &[
|
||||
col!("id", i64),
|
||||
col!("site_id", i64),
|
||||
col!("username", text),
|
||||
col!("access_token", text),
|
||||
col!("api_token", text),
|
||||
col!("balance", f64),
|
||||
col!("balance_used", f64),
|
||||
col!("quota", f64),
|
||||
col!("unit_cost", f64),
|
||||
col!("value_score", f64),
|
||||
col!("status", text),
|
||||
col!("checkin_enabled", bool),
|
||||
col!("last_checkin_at", text),
|
||||
col!("last_balance_refresh", text),
|
||||
col!("extra_config", text),
|
||||
col!("is_pinned", bool),
|
||||
col!("sort_order", i64),
|
||||
col!("oauth_provider", text),
|
||||
col!("oauth_account_key", text),
|
||||
col!("oauth_project_id", text),
|
||||
col!("created_at", text),
|
||||
col!("updated_at", text),
|
||||
],
|
||||
},
|
||||
TableSpec {
|
||||
name: "account_tokens",
|
||||
columns: &[
|
||||
col!("id", i64),
|
||||
col!("account_id", i64),
|
||||
col!("name", text),
|
||||
col!("token", text),
|
||||
col!("source", text),
|
||||
col!("token_group", text),
|
||||
col!("value_status", text),
|
||||
col!("enabled", bool),
|
||||
col!("is_default", bool),
|
||||
col!("created_at", text),
|
||||
col!("updated_at", text),
|
||||
],
|
||||
},
|
||||
TableSpec {
|
||||
name: "downstream_api_keys",
|
||||
columns: &[
|
||||
col!("id", i64),
|
||||
col!("name", text),
|
||||
col!("key", text),
|
||||
col!("description", text),
|
||||
col!("group_name", text),
|
||||
col!("tags", text),
|
||||
col!("enabled", bool),
|
||||
col!("expires_at", text),
|
||||
col!("max_cost", f64),
|
||||
col!("used_cost", f64),
|
||||
col!("max_requests", i64),
|
||||
col!("used_requests", i64),
|
||||
col!("supported_models", text),
|
||||
col!("allowed_route_ids", text),
|
||||
col!("site_weight_multipliers", text),
|
||||
col!("excluded_site_ids", text),
|
||||
col!("excluded_credential_refs", text),
|
||||
col!("last_used_at", text),
|
||||
col!("user_id", i64),
|
||||
col!("created_at", text),
|
||||
col!("updated_at", text),
|
||||
],
|
||||
},
|
||||
TableSpec {
|
||||
name: "token_routes",
|
||||
columns: &[
|
||||
col!("id", i64),
|
||||
col!("model_pattern", text),
|
||||
col!("model_mapping", text),
|
||||
col!("display_name", text),
|
||||
col!("display_icon", text),
|
||||
col!("decision_snapshot", text),
|
||||
col!("decision_refreshed_at", text),
|
||||
col!("routing_strategy", text),
|
||||
col!("route_mode", text),
|
||||
col!("enabled", bool),
|
||||
col!("created_at", text),
|
||||
col!("updated_at", text),
|
||||
],
|
||||
},
|
||||
TableSpec {
|
||||
name: "route_channels",
|
||||
columns: &[
|
||||
col!("id", i64),
|
||||
col!("route_id", i64),
|
||||
col!("account_id", i64),
|
||||
col!("token_id", i64),
|
||||
col!("oauth_route_unit_id", i64),
|
||||
col!("source_model", text),
|
||||
col!("priority", i64),
|
||||
col!("weight", i64),
|
||||
col!("enabled", bool),
|
||||
col!("manual_override", bool),
|
||||
col!("success_count", i64),
|
||||
col!("fail_count", i64),
|
||||
col!("total_latency_ms", i64),
|
||||
col!("total_cost", f64),
|
||||
col!("last_used_at", text),
|
||||
col!("last_selected_at", text),
|
||||
col!("last_fail_at", text),
|
||||
col!("consecutive_fail_count", i64),
|
||||
col!("cooldown_level", i64),
|
||||
col!("cooldown_until", text),
|
||||
],
|
||||
},
|
||||
TableSpec {
|
||||
name: "proxy_logs",
|
||||
columns: &[
|
||||
col!("id", i64),
|
||||
col!("route_id", i64),
|
||||
col!("channel_id", i64),
|
||||
col!("account_id", i64),
|
||||
col!("downstream_api_key_id", i64),
|
||||
col!("model_requested", text),
|
||||
col!("model_actual", text),
|
||||
col!("status", text),
|
||||
col!("http_status", i64),
|
||||
col!("latency_ms", i64),
|
||||
col!("first_byte_latency_ms", i64),
|
||||
col!("prompt_tokens", i64),
|
||||
col!("completion_tokens", i64),
|
||||
col!("total_tokens", i64),
|
||||
col!("estimated_cost", f64),
|
||||
col!("error_message", text),
|
||||
col!("retry_count", i64),
|
||||
col!("is_stream", bool),
|
||||
col!("client_family", text),
|
||||
col!("client_app_id", text),
|
||||
col!("client_app_name", text),
|
||||
col!("client_confidence", text),
|
||||
col!("billing_details", text),
|
||||
col!("created_at", text),
|
||||
],
|
||||
},
|
||||
TableSpec {
|
||||
name: "settings",
|
||||
columns: &[col!("key", text), col!("value", text)],
|
||||
},
|
||||
TableSpec {
|
||||
name: "workspace_assets",
|
||||
columns: &[
|
||||
col!("id", i64),
|
||||
col!("public_id", text),
|
||||
col!("user_id", i64),
|
||||
col!("source", text),
|
||||
col!("kind", text),
|
||||
col!("title", text),
|
||||
col!("prompt", text),
|
||||
col!("model", text),
|
||||
col!("status", text),
|
||||
col!("mime_type", text),
|
||||
col!("file_name", text),
|
||||
col!("file_path", text),
|
||||
col!("remote_url", text),
|
||||
col!("thumbnail_path", text),
|
||||
col!("size_bytes", i64),
|
||||
col!("metadata", text),
|
||||
col!("created_at", text),
|
||||
col!("updated_at", text),
|
||||
col!("deleted_at", text),
|
||||
],
|
||||
},
|
||||
TableSpec {
|
||||
name: "user_announcements",
|
||||
columns: &[
|
||||
col!("id", i64),
|
||||
col!("title", text),
|
||||
col!("content", text),
|
||||
col!("level", text),
|
||||
col!("status", text),
|
||||
col!("audience", text),
|
||||
col!("starts_at", text),
|
||||
col!("ends_at", text),
|
||||
col!("created_at", text),
|
||||
col!("updated_at", text),
|
||||
col!("published_at", text),
|
||||
],
|
||||
},
|
||||
TableSpec {
|
||||
name: "user_announcement_reads",
|
||||
columns: &[
|
||||
col!("announcement_id", i64),
|
||||
col!("user_id", i64),
|
||||
col!("read_at", text),
|
||||
],
|
||||
},
|
||||
];
|
||||
|
||||
pub async fn migrate(
|
||||
sqlite: &SqlitePool,
|
||||
pg: &PgPool,
|
||||
force_truncate: bool,
|
||||
) -> anyhow::Result<MigrationSummary> {
|
||||
ensure_target_is_empty(pg, force_truncate).await?;
|
||||
if force_truncate {
|
||||
truncate_target(pg).await?;
|
||||
}
|
||||
|
||||
let mut summaries = Vec::new();
|
||||
let mut total_rows = 0;
|
||||
for table in TABLES {
|
||||
let rows = copy_table(sqlite, pg, table).await?;
|
||||
summaries.push(TableMigrationSummary {
|
||||
table: table.name,
|
||||
rows,
|
||||
});
|
||||
total_rows += rows;
|
||||
}
|
||||
sync_sequences(pg).await?;
|
||||
Ok(MigrationSummary {
|
||||
tables: summaries,
|
||||
total_rows,
|
||||
force_truncate,
|
||||
})
|
||||
}
|
||||
|
||||
async fn ensure_target_is_empty(pg: &PgPool, force_truncate: bool) -> anyhow::Result<()> {
|
||||
if force_truncate {
|
||||
return Ok(());
|
||||
}
|
||||
let mut non_empty = Vec::new();
|
||||
for table in TABLES {
|
||||
let sql = format!("SELECT COUNT(*) FROM {}", table.name);
|
||||
let count: i64 = sqlx::query_scalar(&sql).fetch_one(pg).await.unwrap_or(0);
|
||||
if count > 0 {
|
||||
non_empty.push(format!("{} ({})", table.name, count));
|
||||
}
|
||||
}
|
||||
if !non_empty.is_empty() {
|
||||
bail!(
|
||||
"target PostgreSQL database is not empty: {}; rerun with --force-truncate to overwrite",
|
||||
non_empty.join(", ")
|
||||
);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn truncate_target(pg: &PgPool) -> anyhow::Result<()> {
|
||||
let table_names = TABLES
|
||||
.iter()
|
||||
.rev()
|
||||
.map(|table| table.name)
|
||||
.collect::<Vec<_>>()
|
||||
.join(", ");
|
||||
let sql = format!("TRUNCATE TABLE {} RESTART IDENTITY CASCADE", table_names);
|
||||
sqlx::query(&sql).execute(pg).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn copy_table(sqlite: &SqlitePool, pg: &PgPool, table: &TableSpec) -> anyhow::Result<i64> {
|
||||
let available = sqlite_columns(sqlite, table.name).await?;
|
||||
let columns = table
|
||||
.columns
|
||||
.iter()
|
||||
.copied()
|
||||
.filter(|column| available.contains_key(column.name))
|
||||
.collect::<Vec<_>>();
|
||||
if columns.is_empty() {
|
||||
return Ok(0);
|
||||
}
|
||||
|
||||
let select_sql = format!(
|
||||
"SELECT {} FROM {} ORDER BY {}",
|
||||
columns
|
||||
.iter()
|
||||
.map(|column| column.name)
|
||||
.collect::<Vec<_>>()
|
||||
.join(", "),
|
||||
table.name,
|
||||
if table.name == "settings" {
|
||||
"key"
|
||||
} else {
|
||||
columns[0].name
|
||||
}
|
||||
);
|
||||
let rows = sqlx::query(&select_sql).fetch_all(sqlite).await?;
|
||||
let insert_sql = format!(
|
||||
"INSERT INTO {} ({}) VALUES ({})",
|
||||
table.name,
|
||||
columns
|
||||
.iter()
|
||||
.map(|column| quote_ident(column.name))
|
||||
.collect::<Vec<_>>()
|
||||
.join(", "),
|
||||
(1..=columns.len())
|
||||
.map(|index| format!("${index}"))
|
||||
.collect::<Vec<_>>()
|
||||
.join(", ")
|
||||
);
|
||||
|
||||
for row in &rows {
|
||||
let mut query = sqlx::query(&insert_sql);
|
||||
for column in &columns {
|
||||
query = bind_column(query, row, column)?;
|
||||
}
|
||||
query
|
||||
.execute(pg)
|
||||
.await
|
||||
.with_context(|| format!("insert into {}", table.name))?;
|
||||
}
|
||||
Ok(rows.len() as i64)
|
||||
}
|
||||
|
||||
fn bind_column<'q>(
|
||||
query: sqlx::query::Query<'q, sqlx::Postgres, sqlx::postgres::PgArguments>,
|
||||
row: &SqliteRow,
|
||||
column: &ColumnSpec,
|
||||
) -> anyhow::Result<sqlx::query::Query<'q, sqlx::Postgres, sqlx::postgres::PgArguments>> {
|
||||
Ok(match column.kind {
|
||||
ColumnKind::I64 => query.bind(row.try_get::<Option<i64>, _>(column.name)?),
|
||||
ColumnKind::F64 => query.bind(row.try_get::<Option<f64>, _>(column.name)?),
|
||||
ColumnKind::Bool => query.bind(row.try_get::<Option<bool>, _>(column.name)?),
|
||||
ColumnKind::Text => query.bind(row.try_get::<Option<String>, _>(column.name)?),
|
||||
})
|
||||
}
|
||||
|
||||
async fn sqlite_columns(sqlite: &SqlitePool, table: &str) -> anyhow::Result<HashMap<String, ()>> {
|
||||
let rows = sqlx::query(&format!("PRAGMA table_info({table})"))
|
||||
.fetch_all(sqlite)
|
||||
.await?;
|
||||
Ok(rows
|
||||
.into_iter()
|
||||
.filter_map(|row| row.try_get::<String, _>("name").ok())
|
||||
.map(|name| (name, ()))
|
||||
.collect())
|
||||
}
|
||||
|
||||
async fn sync_sequences(pg: &PgPool) -> anyhow::Result<()> {
|
||||
for table in TABLES
|
||||
.iter()
|
||||
.filter(|table| table.columns.iter().any(|column| column.name == "id"))
|
||||
{
|
||||
let sql = format!(
|
||||
"SELECT setval(pg_get_serial_sequence('{table}', 'id'), COALESCE((SELECT MAX(id) FROM {table}), 1), (SELECT COUNT(*) FROM {table}) > 0)",
|
||||
table = table.name
|
||||
);
|
||||
let _ = sqlx::query(&sql).execute(pg).await;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn quote_ident(name: &str) -> String {
|
||||
if name == "read" {
|
||||
"\"read\"".to_string()
|
||||
} else {
|
||||
name.to_string()
|
||||
}
|
||||
}
|
||||
+17
-1
@@ -24,6 +24,12 @@ pub enum AppError {
|
||||
#[error("Rate limited")]
|
||||
RateLimited,
|
||||
|
||||
#[error("Rate limited: {message}")]
|
||||
RateLimitedWithRetry {
|
||||
message: String,
|
||||
retry_after_secs: u64,
|
||||
},
|
||||
|
||||
#[error("Service unavailable: {0}")]
|
||||
ServiceUnavailable(String),
|
||||
|
||||
@@ -46,18 +52,28 @@ impl IntoResponse for AppError {
|
||||
AppError::BadRequest(msg) => (StatusCode::BAD_REQUEST, msg.clone()),
|
||||
AppError::UpstreamError(msg) => (StatusCode::BAD_GATEWAY, msg.clone()),
|
||||
AppError::RateLimited => (StatusCode::TOO_MANY_REQUESTS, "rate_limited".into()),
|
||||
AppError::RateLimitedWithRetry { message, .. } => {
|
||||
(StatusCode::TOO_MANY_REQUESTS, message.clone())
|
||||
}
|
||||
AppError::ServiceUnavailable(msg) => (StatusCode::SERVICE_UNAVAILABLE, msg.clone()),
|
||||
AppError::GatewayTimeout(msg) => (StatusCode::GATEWAY_TIMEOUT, msg.clone()),
|
||||
AppError::Internal(e) => (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()),
|
||||
AppError::Database(e) => (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()),
|
||||
};
|
||||
let retry_after_secs = match &self {
|
||||
AppError::RateLimitedWithRetry {
|
||||
retry_after_secs, ..
|
||||
} => Some(*retry_after_secs),
|
||||
_ => None,
|
||||
};
|
||||
(
|
||||
status,
|
||||
Json(json!({
|
||||
"error": {
|
||||
"message": message,
|
||||
"type": status.canonical_reason(),
|
||||
}
|
||||
},
|
||||
"retry_after_secs": retry_after_secs,
|
||||
})),
|
||||
)
|
||||
.into_response()
|
||||
|
||||
+5
-11
@@ -7,6 +7,7 @@ pub mod routes;
|
||||
pub mod services;
|
||||
pub mod shared;
|
||||
pub mod transformers;
|
||||
pub mod workflow;
|
||||
|
||||
use sqlx::SqlitePool;
|
||||
use std::sync::Arc;
|
||||
@@ -16,8 +17,8 @@ pub struct AppState {
|
||||
pub config: config::AppConfig,
|
||||
pub db: Arc<SqlitePool>,
|
||||
pub comfyui: Arc<services::comfyui_process::ComfyUIProcess>,
|
||||
pub comfyui_pool: Arc<services::comfyui_pool::ComfyUIPool>,
|
||||
pub builtin_keys: Arc<services::builtin_keys::BuiltinKeys>,
|
||||
pub registration_limiter: Arc<services::registration_limiter::RegistrationLimiter>,
|
||||
}
|
||||
|
||||
impl AppState {
|
||||
@@ -40,21 +41,14 @@ impl AppState {
|
||||
comfyui_agent_secret.clone(),
|
||||
);
|
||||
|
||||
// Per-user multi-tenant ComfyUI pool. Backed by Docker today; the
|
||||
// ContainerRuntime trait makes K3S a drop-in v2.
|
||||
let docker = Arc::new(services::docker_runtime::DockerRuntime::new());
|
||||
let pool_cfg = services::comfyui_pool::PoolConfig::from_env(
|
||||
format!("{}/v1", boosapi_url.trim_end_matches('/')),
|
||||
comfyui_agent_secret,
|
||||
);
|
||||
let comfyui_pool = services::comfyui_pool::ComfyUIPool::new(docker, pool_cfg);
|
||||
|
||||
Self {
|
||||
config,
|
||||
db: Arc::new(db),
|
||||
comfyui,
|
||||
comfyui_pool,
|
||||
builtin_keys: Arc::new(builtin_keys),
|
||||
registration_limiter: Arc::new(
|
||||
services::registration_limiter::RegistrationLimiter::new(),
|
||||
),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
+40
-24
@@ -1,7 +1,7 @@
|
||||
use anyhow::Context;
|
||||
use axum::serve;
|
||||
use boosapi::{config::AppConfig, db, services::builtin_keys, AppState};
|
||||
use clap::Parser;
|
||||
use clap::{Parser, Subcommand};
|
||||
use std::ffi::OsString;
|
||||
use std::net::SocketAddr;
|
||||
use tokio::signal;
|
||||
@@ -15,6 +15,9 @@ use tracing_subscriber::EnvFilter;
|
||||
about = "AI API Gateway — Rust rewrite of BoosAPI"
|
||||
)]
|
||||
struct Cli {
|
||||
#[command(subcommand)]
|
||||
command: Option<Commands>,
|
||||
|
||||
/// Config file path
|
||||
#[arg(short, long, default_value = "config.toml")]
|
||||
config: String,
|
||||
@@ -40,10 +43,42 @@ struct Cli {
|
||||
check_config: bool,
|
||||
}
|
||||
|
||||
#[derive(Subcommand, Debug)]
|
||||
enum Commands {
|
||||
/// Copy core data from SQLite into PostgreSQL and exit
|
||||
MigrateSqliteToPostgres {
|
||||
/// Source SQLite URL, e.g. sqlite:boosapi.db?mode=rwc
|
||||
#[arg(long)]
|
||||
sqlite_url: String,
|
||||
|
||||
/// Target PostgreSQL URL, e.g. postgres://user:pass@host:5432/boosapi
|
||||
#[arg(long)]
|
||||
postgres_url: String,
|
||||
|
||||
/// Truncate target business tables before importing
|
||||
#[arg(long)]
|
||||
force_truncate: bool,
|
||||
},
|
||||
}
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> anyhow::Result<()> {
|
||||
let args = Cli::parse_from(normalize_legacy_args());
|
||||
|
||||
if let Some(Commands::MigrateSqliteToPostgres {
|
||||
sqlite_url,
|
||||
postgres_url,
|
||||
force_truncate,
|
||||
}) = &args.command
|
||||
{
|
||||
let sqlite = db::create_pool(sqlite_url, 1).await?;
|
||||
let pg = db::create_pg_pool(postgres_url, 2).await?;
|
||||
db::run_migrations(db::DbKind::Postgres, None, Some(&pg)).await?;
|
||||
let summary = db::sqlite_to_postgres::migrate(&sqlite, &pg, *force_truncate).await?;
|
||||
println!("{}", serde_json::to_string_pretty(&summary)?);
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
// Load config
|
||||
let mut config = AppConfig::load_from_file(&args.config)?;
|
||||
|
||||
@@ -73,12 +108,14 @@ async fn main() -> anyhow::Result<()> {
|
||||
|
||||
tracing::info!("Loaded configuration from {}", args.config);
|
||||
|
||||
let db_kind = db::DbKind::from_url(&config.db.url)?;
|
||||
|
||||
// Create database pool
|
||||
let pool = db::create_pool(&config.db.url, config.db.max_connections).await?;
|
||||
tracing::info!("Database connected");
|
||||
|
||||
// Run migrations
|
||||
sqlx::migrate!("./migrations").run(&pool).await?;
|
||||
db::run_migrations(db_kind, Some(&pool), None).await?;
|
||||
tracing::info!("Database migrations applied");
|
||||
|
||||
if args.migrate_only {
|
||||
@@ -105,20 +142,6 @@ async fn main() -> anyhow::Result<()> {
|
||||
}
|
||||
let comfyui_process = state.comfyui.clone();
|
||||
|
||||
// Multi-tenant ComfyUI pool: adopt any containers that survived a
|
||||
// previous API restart, then start the idle-reaper task. Only
|
||||
// wired in when the `multi_tenant_comfy` feature is on; under the
|
||||
// legacy singleton path the pool is constructed but unused, and
|
||||
// `comfyui_pool` is `None` so we skip shutdown_all on exit.
|
||||
#[cfg(feature = "multi_tenant_comfy")]
|
||||
let comfyui_pool = {
|
||||
state.comfyui_pool.reconcile_on_startup().await;
|
||||
boosapi::services::comfyui_pool::spawn_reaper(state.comfyui_pool.clone());
|
||||
Some(state.comfyui_pool.clone())
|
||||
};
|
||||
#[cfg(not(feature = "multi_tenant_comfy"))]
|
||||
let comfyui_pool: Option<std::sync::Arc<boosapi::services::comfyui_pool::ComfyUIPool>> = None;
|
||||
|
||||
// Build router
|
||||
let app = boosapi::routes::build_router(state);
|
||||
|
||||
@@ -128,7 +151,7 @@ async fn main() -> anyhow::Result<()> {
|
||||
|
||||
serve(
|
||||
tokio::net::TcpListener::bind(addr).await?,
|
||||
app.into_make_service(),
|
||||
app.into_make_service_with_connect_info::<SocketAddr>(),
|
||||
)
|
||||
.with_graceful_shutdown(shutdown_signal())
|
||||
.await?;
|
||||
@@ -137,13 +160,6 @@ async fn main() -> anyhow::Result<()> {
|
||||
tracing::warn!(%err, "ComfyUI embedded process did not stop cleanly");
|
||||
}
|
||||
|
||||
// Tear down any per-user ComfyUI containers we spawned. This is
|
||||
// best-effort — a hard kill of the API leaves them running and the
|
||||
// next startup will adopt them via reconcile_on_startup().
|
||||
if let Some(pool) = comfyui_pool {
|
||||
pool.shutdown_all().await;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
|
||||
@@ -267,6 +267,7 @@ mod tests {
|
||||
route_id: Some(10),
|
||||
site_id: 20,
|
||||
account_id: Some(30),
|
||||
account_unit_cost: None,
|
||||
base_url: "https://example.test/v1".into(),
|
||||
api_key: "key".into(),
|
||||
model: "gpt-5.5".into(),
|
||||
|
||||
@@ -11,6 +11,7 @@ pub struct SelectedChannel {
|
||||
pub route_id: Option<i64>,
|
||||
pub site_id: i64,
|
||||
pub account_id: Option<i64>,
|
||||
pub account_unit_cost: Option<f64>,
|
||||
pub base_url: String,
|
||||
pub api_key: String,
|
||||
pub model: String,
|
||||
|
||||
+28
-24
@@ -76,6 +76,23 @@ fn looks_like_sse(text: &str) -> bool {
|
||||
text.trim_start().starts_with("data:")
|
||||
}
|
||||
|
||||
fn openai_compatible_base_url(base_url: &str) -> String {
|
||||
let trimmed = base_url.trim().trim_end_matches('/');
|
||||
if trimmed.ends_with("/v1") || trimmed.contains("/v1/") {
|
||||
trimmed.to_string()
|
||||
} else {
|
||||
format!("{trimmed}/v1")
|
||||
}
|
||||
}
|
||||
|
||||
fn upstream_url(base_url: &str, path_suffix: &str) -> String {
|
||||
format!(
|
||||
"{}/{}",
|
||||
openai_compatible_base_url(base_url),
|
||||
path_suffix.trim_start_matches('/')
|
||||
)
|
||||
}
|
||||
|
||||
fn aggregate_chat_sse(text: &str, fallback_model: &str) -> Option<Value> {
|
||||
let mut id: Option<String> = None;
|
||||
let mut created: Option<i64> = None;
|
||||
@@ -162,7 +179,7 @@ pub async fn execute_chat_completions(
|
||||
body: Value,
|
||||
) -> AppResult<Value> {
|
||||
let client = http_client();
|
||||
let upstream_url = format!("{}/chat/completions", route.base_url.trim_end_matches('/'));
|
||||
let upstream_url = upstream_url(&route.base_url, "/chat/completions");
|
||||
|
||||
let resp = client
|
||||
.post(&upstream_url)
|
||||
@@ -279,8 +296,8 @@ pub async fn execute_anthropic_messages_stream(
|
||||
let mut stream = resp.bytes_stream();
|
||||
use futures::StreamExt;
|
||||
while let Some(chunk_result) = stream.next().await {
|
||||
let bytes = chunk_result
|
||||
.map_err(|e| AppError::UpstreamError(format!("Stream error: {e}")))?;
|
||||
let bytes =
|
||||
chunk_result.map_err(|e| AppError::UpstreamError(format!("Stream error: {e}")))?;
|
||||
let text = String::from_utf8_lossy(&bytes).to_string();
|
||||
frames.extend(normalize_anthropic_sse_frames(&text));
|
||||
}
|
||||
@@ -328,7 +345,9 @@ fn normalize_anthropic_sse_frame(frame: &str) -> Option<String> {
|
||||
|
||||
let choice = &event.choices[0];
|
||||
if choice.delta.role.is_some() {
|
||||
let id = event.id.unwrap_or_else(|| format!("msg_{}", uuid::Uuid::new_v4().simple()));
|
||||
let id = event
|
||||
.id
|
||||
.unwrap_or_else(|| format!("msg_{}", uuid::Uuid::new_v4().simple()));
|
||||
return Some(format!(
|
||||
"event: message_start\ndata: {}\n\n",
|
||||
serde_json::json!({
|
||||
@@ -392,11 +411,7 @@ pub async fn execute_streaming_request(
|
||||
body: Value,
|
||||
) -> AppResult<Vec<String>> {
|
||||
let client = http_client();
|
||||
let upstream_url = format!(
|
||||
"{}/{}",
|
||||
route.base_url.trim_end_matches('/'),
|
||||
path_suffix.trim_start_matches('/')
|
||||
);
|
||||
let upstream_url = upstream_url(&route.base_url, path_suffix);
|
||||
|
||||
let resp = client
|
||||
.post(&upstream_url)
|
||||
@@ -448,11 +463,7 @@ pub async fn execute_proxy_request(
|
||||
body: Value,
|
||||
) -> AppResult<Value> {
|
||||
let client = http_client();
|
||||
let upstream_url = format!(
|
||||
"{}/{}",
|
||||
route.base_url.trim_end_matches('/'),
|
||||
path_suffix.trim_start_matches('/')
|
||||
);
|
||||
let upstream_url = upstream_url(&route.base_url, path_suffix);
|
||||
|
||||
let resp = client
|
||||
.post(&upstream_url)
|
||||
@@ -483,11 +494,7 @@ pub async fn execute_proxy_request(
|
||||
/// Generic GET proxy request to an upstream endpoint path.
|
||||
pub async fn execute_proxy_get(route: UpstreamRoute, path_suffix: &str) -> AppResult<Value> {
|
||||
let client = http_client();
|
||||
let upstream_url = format!(
|
||||
"{}/{}",
|
||||
route.base_url.trim_end_matches('/'),
|
||||
path_suffix.trim_start_matches('/')
|
||||
);
|
||||
let upstream_url = upstream_url(&route.base_url, path_suffix);
|
||||
|
||||
let resp = client
|
||||
.get(&upstream_url)
|
||||
@@ -602,10 +609,7 @@ pub async fn execute_image_generation(
|
||||
body: Value,
|
||||
) -> AppResult<Value> {
|
||||
let client = http_client();
|
||||
let upstream_url = format!(
|
||||
"{}/images/generations",
|
||||
route.base_url.trim_end_matches('/')
|
||||
);
|
||||
let upstream_url = upstream_url(&route.base_url, "/images/generations");
|
||||
|
||||
let resp = client
|
||||
.post(&upstream_url)
|
||||
@@ -664,7 +668,7 @@ pub async fn execute_image_edit(
|
||||
use reqwest::multipart::{Form, Part};
|
||||
|
||||
let client = http_client();
|
||||
let upstream_url = format!("{}/images/edits", route.base_url.trim_end_matches('/'));
|
||||
let upstream_url = upstream_url(&route.base_url, "/images/edits");
|
||||
|
||||
let mut form = Form::new();
|
||||
for field in fields {
|
||||
|
||||
@@ -29,6 +29,7 @@ struct ChannelRow {
|
||||
site_api_key: Option<String>,
|
||||
account_access_token: String,
|
||||
account_api_token: Option<String>,
|
||||
account_unit_cost: Option<f64>,
|
||||
token_name: Option<String>,
|
||||
token_value: Option<String>,
|
||||
source_model: Option<String>,
|
||||
@@ -148,6 +149,7 @@ impl TokenRouter {
|
||||
s.api_key AS site_api_key,
|
||||
a.access_token AS account_access_token,
|
||||
a.api_token AS account_api_token,
|
||||
a.unit_cost AS account_unit_cost,
|
||||
at.name AS token_name,
|
||||
at.token AS token_value,
|
||||
rc.source_model AS source_model
|
||||
@@ -198,6 +200,7 @@ impl TokenRouter {
|
||||
route_id: Some(row.route_id),
|
||||
site_id: row.site_id,
|
||||
account_id: Some(row.account_id),
|
||||
account_unit_cost: row.account_unit_cost,
|
||||
base_url: row.base_url,
|
||||
api_key: api_key
|
||||
.trim()
|
||||
|
||||
@@ -100,6 +100,14 @@ struct IdPath {
|
||||
id: i64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct DatabaseConnectionRequest {
|
||||
dialect: Option<String>,
|
||||
connection_string: Option<String>,
|
||||
ssl: Option<bool>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct RouteChannelPath {
|
||||
route_id: i64,
|
||||
@@ -599,23 +607,102 @@ async fn test_system_proxy() -> Json<Value> {
|
||||
}
|
||||
|
||||
async fn database_runtime(State(state): State<AppState>) -> Json<Value> {
|
||||
let dialect = crate::db::DbKind::from_url(&state.config.db.url)
|
||||
.map(|kind| kind.as_str().to_string())
|
||||
.unwrap_or_else(|_| "unknown".to_string());
|
||||
Json(json!({
|
||||
"dialect": "sqlite",
|
||||
"connectionString": state.config.db.url,
|
||||
"dialect": dialect,
|
||||
"connectionString": mask_connection_string(&state.config.db.url),
|
||||
"ssl": state.config.db.ssl,
|
||||
"active": {
|
||||
"dialect": dialect,
|
||||
"connection": mask_connection_string(&state.config.db.url),
|
||||
"ssl": state.config.db.ssl,
|
||||
},
|
||||
"saved": state.config.db.runtime_db_url.as_ref().map(|url| json!({
|
||||
"dialect": state.config.db.runtime_db_dialect.as_deref().unwrap_or("postgres"),
|
||||
"connection": mask_connection_string(url),
|
||||
"ssl": state.config.db.runtime_db_ssl.unwrap_or(false),
|
||||
})),
|
||||
"restartRequired": false,
|
||||
}))
|
||||
}
|
||||
|
||||
async fn update_database_runtime() -> Json<Value> {
|
||||
Json(json!({ "success": true, "message": "database config is managed by config.toml" }))
|
||||
Json(json!({
|
||||
"success": false,
|
||||
"message": "数据库运行配置由 config.toml / BOOSAPI_DB__URL 管理；修改后需要重启服务",
|
||||
"restartRequired": true,
|
||||
}))
|
||||
}
|
||||
|
||||
async fn test_database_connection() -> Json<Value> {
|
||||
Json(json!({ "success": true }))
|
||||
async fn test_database_connection(
|
||||
Json(req): Json<DatabaseConnectionRequest>,
|
||||
) -> AppResult<Json<Value>> {
|
||||
let connection = req
|
||||
.connection_string
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.ok_or_else(|| AppError::BadRequest("connectionString is required".into()))?;
|
||||
let detected = crate::db::test_connection(connection, 2)
|
||||
.await
|
||||
.map_err(|err| AppError::BadRequest(err.to_string()))?;
|
||||
if let Some(expected) = req
|
||||
.dialect
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
{
|
||||
if expected != detected.as_str() {
|
||||
return Err(AppError::BadRequest(format!(
|
||||
"dialect mismatch: selected {}, detected {}",
|
||||
expected,
|
||||
detected.as_str()
|
||||
)));
|
||||
}
|
||||
}
|
||||
Ok(Json(json!({
|
||||
"success": true,
|
||||
"dialect": detected.as_str(),
|
||||
"connection": mask_connection_string(connection),
|
||||
"ssl": req.ssl.unwrap_or(false),
|
||||
})))
|
||||
}
|
||||
|
||||
async fn database_migrate() -> Json<Value> {
|
||||
Json(json!({ "success": true }))
|
||||
Json(json!({
|
||||
"success": false,
|
||||
"message": "Web 迁移入口已禁用。请在服务器执行 boosapi migrate-sqlite-to-postgres --sqlite-url <sqlite-url> --postgres-url <postgres-url>，需要覆盖目标库时追加 --force-truncate。",
|
||||
"cliOnly": true,
|
||||
}))
|
||||
}
|
||||
|
||||
fn mask_connection_string(raw: &str) -> String {
|
||||
let trimmed = raw.trim();
|
||||
if trimmed.starts_with("sqlite:") {
|
||||
return trimmed.to_string();
|
||||
}
|
||||
let Some(scheme_end) = trimmed.find("://") else {
|
||||
return trimmed.to_string();
|
||||
};
|
||||
let after_scheme = scheme_end + 3;
|
||||
let Some(at_rel) = trimmed[after_scheme..].find('@') else {
|
||||
return trimmed.to_string();
|
||||
};
|
||||
let at = after_scheme + at_rel;
|
||||
let credentials = &trimmed[after_scheme..at];
|
||||
let masked_credentials = if let Some(colon) = credentials.find(':') {
|
||||
format!("{}:****", &credentials[..colon])
|
||||
} else {
|
||||
"****".to_string()
|
||||
};
|
||||
format!(
|
||||
"{}{}{}",
|
||||
&trimmed[..after_scheme],
|
||||
masked_credentials,
|
||||
&trimmed[at..]
|
||||
)
|
||||
}
|
||||
|
||||
async fn export_backup(
|
||||
|
||||
@@ -19,7 +19,7 @@ mod accounts;
|
||||
mod compat;
|
||||
mod downstream_keys;
|
||||
mod model_candidates;
|
||||
mod models_marketplace;
|
||||
pub(crate) mod models_marketplace;
|
||||
mod sites;
|
||||
mod tokens;
|
||||
mod users;
|
||||
@@ -63,6 +63,7 @@ pub fn routes(state: AppState) -> Router<AppState> {
|
||||
"/api/sites",
|
||||
get(sites::list_sites).post(sites::create_site),
|
||||
)
|
||||
.route("/api/sites/detect", post(sites::detect_site))
|
||||
.route(
|
||||
"/api/sites/{id}",
|
||||
get(sites::get_site)
|
||||
|
||||
@@ -12,6 +12,11 @@ use serde::{Deserialize, Serialize};
|
||||
use serde_json::{json, Value};
|
||||
use tracing::info;
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct DetectSiteRequest {
|
||||
pub url: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct CreateSiteRequest {
|
||||
pub name: String,
|
||||
@@ -90,6 +95,18 @@ struct SiteApiEndpointRow {
|
||||
last_failure_reason: Option<String>,
|
||||
}
|
||||
|
||||
/// POST /api/sites/detect
|
||||
pub async fn detect_site(Json(payload): Json<DetectSiteRequest>) -> AppResult<Json<Value>> {
|
||||
let url = normalize_site_url(&payload.url)?;
|
||||
let (platform, initialization_preset_id) = detect_site_platform(&url);
|
||||
|
||||
Ok(Json(json!({
|
||||
"url": url,
|
||||
"platform": platform,
|
||||
"initializationPresetId": initialization_preset_id,
|
||||
})))
|
||||
}
|
||||
|
||||
/// GET /api/sites
|
||||
pub async fn list_sites(State(state): State<AppState>) -> AppResult<Json<Value>> {
|
||||
info!("Listing all sites");
|
||||
@@ -435,6 +452,102 @@ pub async fn update_site_disabled_models(
|
||||
Ok(Json(json!({ "models": normalised })))
|
||||
}
|
||||
|
||||
fn normalize_site_url(raw: &str) -> AppResult<String> {
|
||||
let trimmed = raw.trim().trim_end_matches('/').to_string();
|
||||
if trimmed.is_empty() {
|
||||
return Err(AppError::BadRequest("url is required".into()));
|
||||
}
|
||||
let lower = trimmed.to_ascii_lowercase();
|
||||
if !(lower.starts_with("http://") || lower.starts_with("https://")) {
|
||||
return Err(AppError::BadRequest(
|
||||
"url must start with http:// or https://".into(),
|
||||
));
|
||||
}
|
||||
Ok(trimmed)
|
||||
}
|
||||
|
||||
fn detect_site_platform(url: &str) -> (&'static str, Option<&'static str>) {
|
||||
let lower = url.to_ascii_lowercase();
|
||||
|
||||
if lower.starts_with("https://coding.dashscope.aliyuncs.com/apps/anthropic") {
|
||||
return ("claude", Some("codingplan-claude"));
|
||||
}
|
||||
if lower.starts_with("https://coding.dashscope.aliyuncs.com/v1") {
|
||||
return ("openai", Some("codingplan-openai"));
|
||||
}
|
||||
if lower.starts_with("https://open.bigmodel.cn/api/anthropic") {
|
||||
return ("claude", Some("zhipu-coding-plan-claude"));
|
||||
}
|
||||
if lower.starts_with("https://open.bigmodel.cn/api/coding/paas/v4") {
|
||||
return ("openai", Some("zhipu-coding-plan-openai"));
|
||||
}
|
||||
if lower.starts_with("https://api.deepseek.com/anthropic") {
|
||||
return ("claude", Some("deepseek-claude"));
|
||||
}
|
||||
if lower.starts_with("https://api.deepseek.com/v1") || lower == "https://api.deepseek.com" {
|
||||
return ("openai", Some("deepseek-openai"));
|
||||
}
|
||||
if lower.starts_with("https://api.moonshot.cn/anthropic") {
|
||||
return ("claude", Some("moonshot-claude"));
|
||||
}
|
||||
if lower.starts_with("https://api.moonshot.cn/v1") {
|
||||
return ("openai", Some("moonshot-openai"));
|
||||
}
|
||||
if lower.starts_with("https://api.minimaxi.com/anthropic") {
|
||||
return ("claude", Some("minimax-claude"));
|
||||
}
|
||||
if lower.starts_with("https://api.minimaxi.com/v1") {
|
||||
return ("openai", Some("minimax-openai"));
|
||||
}
|
||||
if lower.starts_with("https://api-inference.modelscope.cn/v1") {
|
||||
return ("openai", Some("modelscope-openai"));
|
||||
}
|
||||
if lower.starts_with("https://api-inference.modelscope.cn") {
|
||||
return ("claude", Some("modelscope-claude"));
|
||||
}
|
||||
if lower.starts_with("https://ark.cn-beijing.volces.com/api/coding/v3") {
|
||||
return ("openai", Some("doubao-coding-openai"));
|
||||
}
|
||||
if lower.starts_with("https://ark.cn-beijing.volces.com/api/v3")
|
||||
|| lower.starts_with("https://ark.cn-beijing.volces.com")
|
||||
{
|
||||
return ("openai", Some("volcengine-ark-openai"));
|
||||
}
|
||||
|
||||
if lower.contains("anthropic") || lower.contains("claude") {
|
||||
return ("claude", None);
|
||||
}
|
||||
if lower.contains("gemini") || lower.contains("googleapis.com") {
|
||||
return ("gemini", None);
|
||||
}
|
||||
if lower.contains("new-api") {
|
||||
return ("new-api", None);
|
||||
}
|
||||
if lower.contains("one-api") {
|
||||
return ("one-api", None);
|
||||
}
|
||||
if lower.contains("anyrouter") {
|
||||
return ("anyrouter", None);
|
||||
}
|
||||
if lower.contains("veloera") {
|
||||
return ("veloera", None);
|
||||
}
|
||||
if lower.contains("one-hub") {
|
||||
return ("one-hub", None);
|
||||
}
|
||||
if lower.contains("done-hub") {
|
||||
return ("done-hub", None);
|
||||
}
|
||||
if lower.contains("sub2api") {
|
||||
return ("sub2api", None);
|
||||
}
|
||||
if lower.contains("cliproxy") || lower.contains("cpa") {
|
||||
return ("cliproxyapi", None);
|
||||
}
|
||||
|
||||
("openai", None)
|
||||
}
|
||||
|
||||
async fn replace_site_api_endpoints(
|
||||
tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>,
|
||||
site_id: i64,
|
||||
|
||||
+40
-4
@@ -8,13 +8,14 @@ use crate::{
|
||||
AppState,
|
||||
};
|
||||
use axum::{
|
||||
extract::{Path, State},
|
||||
extract::{ConnectInfo, Path, State},
|
||||
http::HeaderMap,
|
||||
routing::{get, post},
|
||||
Json, Router,
|
||||
};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::{json, Value};
|
||||
use std::net::SocketAddr;
|
||||
use tracing::info;
|
||||
|
||||
// ── Request / Response types ─────────────────────────────────────────
|
||||
@@ -75,6 +76,7 @@ pub fn routes() -> Router<AppState> {
|
||||
/// POST /api/auth/register
|
||||
async fn register(
|
||||
State(state): State<AppState>,
|
||||
ConnectInfo(addr): ConnectInfo<SocketAddr>,
|
||||
Json(payload): Json<RegisterRequest>,
|
||||
) -> AppResult<Json<Value>> {
|
||||
info!("Registering user: {}", payload.email);
|
||||
@@ -90,6 +92,23 @@ async fn register(
|
||||
));
|
||||
}
|
||||
|
||||
let client_ip = Some(addr.ip());
|
||||
if let Err(limit) = state
|
||||
.registration_limiter
|
||||
.check_and_record_attempt(client_ip, &payload.email)
|
||||
{
|
||||
tracing::warn!(
|
||||
email = %payload.email,
|
||||
client_ip = ?client_ip,
|
||||
retry_after_secs = limit.retry_after_secs,
|
||||
"registration rate limited",
|
||||
);
|
||||
return Err(AppError::RateLimitedWithRetry {
|
||||
message: limit.message.to_string(),
|
||||
retry_after_secs: limit.retry_after_secs,
|
||||
});
|
||||
}
|
||||
|
||||
// Check if email already exists
|
||||
let existing: Option<User> = sqlx::query_as::<_, User>("SELECT * FROM users WHERE email = ?")
|
||||
.bind(&payload.email)
|
||||
@@ -103,17 +122,33 @@ async fn register(
|
||||
let password_hash = crate::auth::password::hash_password(&payload.password)
|
||||
.map_err(|e| AppError::Internal(anyhow::anyhow!("Failed to hash password: {e}")))?;
|
||||
|
||||
let mut tx = state.db.begin().await?;
|
||||
|
||||
let result = sqlx::query(
|
||||
"INSERT INTO users (username, email, password_hash, role) VALUES (?, ?, ?, 'user')",
|
||||
)
|
||||
.bind(&payload.username)
|
||||
.bind(&payload.email)
|
||||
.bind(&password_hash)
|
||||
.execute(&*state.db)
|
||||
.execute(&mut *tx)
|
||||
.await?;
|
||||
|
||||
let user_id = result.last_insert_rowid();
|
||||
let default_key = format!("sk-user-{}", uuid::Uuid::new_v4().simple());
|
||||
let default_key_name = format!("{} 的默认密钥", payload.username.trim());
|
||||
sqlx::query(
|
||||
"INSERT INTO downstream_api_keys (name, key, user_id, enabled) VALUES (?, ?, ?, 1)",
|
||||
)
|
||||
.bind(&default_key_name)
|
||||
.bind(&default_key)
|
||||
.bind(user_id)
|
||||
.execute(&mut *tx)
|
||||
.await?;
|
||||
tx.commit().await?;
|
||||
state.registration_limiter.record_success(client_ip);
|
||||
|
||||
let user: User = sqlx::query_as::<_, User>("SELECT * FROM users WHERE id = ?")
|
||||
.bind(result.last_insert_rowid())
|
||||
.bind(user_id)
|
||||
.fetch_one(&*state.db)
|
||||
.await?;
|
||||
|
||||
@@ -144,7 +179,8 @@ async fn register(
|
||||
"role": user.role,
|
||||
"status": user.status,
|
||||
"createdAt": user.created_at,
|
||||
}
|
||||
},
|
||||
"defaultApiKey": default_key
|
||||
}
|
||||
})))
|
||||
}
|
||||
|
||||
@@ -1,457 +0,0 @@
|
||||
//! `/comfy/*` reverse proxy.
|
||||
//!
|
||||
//! Authenticates the caller via JWT/admin token (Authorization header, cookie,
|
||||
//! or `?token=` query param where browsers can't set headers),
|
||||
//! looks up the per-user container in [`ComfyUIPool`], and forwards
|
||||
//! both HTTP and WS traffic to that container's loopback port.
|
||||
//!
|
||||
//! The proxy strips `/comfy` on requests so ComfyUI can keep serving from
|
||||
//! its normal root while browsers see a scoped iframe URL.
|
||||
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use axum::{
|
||||
body::Body,
|
||||
extract::{
|
||||
ws::{Message, WebSocket, WebSocketUpgrade},
|
||||
OriginalUri, RawQuery, State,
|
||||
},
|
||||
http::{
|
||||
header::{
|
||||
CONNECTION, HOST, PROXY_AUTHENTICATE, PROXY_AUTHORIZATION, TE, TRAILER,
|
||||
TRANSFER_ENCODING, UPGRADE,
|
||||
},
|
||||
HeaderMap, HeaderName, HeaderValue, Method, Request, Response, StatusCode, Uri,
|
||||
},
|
||||
response::IntoResponse,
|
||||
routing::{any, get},
|
||||
Router,
|
||||
};
|
||||
use futures::{SinkExt, StreamExt};
|
||||
use tokio_tungstenite::tungstenite::Message as TMessage;
|
||||
use tracing::{debug, warn};
|
||||
|
||||
use crate::services::comfyui_pool::{ComfyUIPool, ContainerState, PoolError};
|
||||
use crate::AppState;
|
||||
|
||||
const COLD_START_HOLD: Duration = Duration::from_secs(30);
|
||||
const COLD_START_POLL_INTERVAL: Duration = Duration::from_millis(500);
|
||||
const COMFY_SESSION_COOKIE: &str = "boosapi_comfy_token";
|
||||
const ADMIN_WORKSPACE_USER_ID: i64 = 0;
|
||||
|
||||
pub fn routes() -> Router<AppState> {
|
||||
Router::new()
|
||||
// WS upgrade endpoints — must be matched before the catch-all.
|
||||
.route("/comfy/ws", get(ws_handler))
|
||||
.route("/comfy/ws/", get(ws_handler))
|
||||
// Catch-all for HTTP. `{*path}` is axum 0.8 wildcard syntax.
|
||||
.route("/comfy", any(http_handler))
|
||||
.route("/comfy/", any(http_handler))
|
||||
.route("/comfy/{*path}", any(http_handler))
|
||||
}
|
||||
|
||||
// ── Auth ───────────────────────────────────────────────────────────────
|
||||
|
||||
fn user_id_from_token(state: &AppState, token: &str) -> Option<i64> {
|
||||
let token = token.trim();
|
||||
if token.is_empty() {
|
||||
return None;
|
||||
}
|
||||
if token == state.config.auth.auth_token || token == state.config.auth.proxy_token {
|
||||
return Some(ADMIN_WORKSPACE_USER_ID);
|
||||
}
|
||||
let claims = crate::auth::jwt::validate_jwt(&state.config.auth.jwt_secret, token).ok()?;
|
||||
claims.sub.parse::<i64>().ok()
|
||||
}
|
||||
|
||||
/// Resolve the calling user from `Authorization: Bearer <token>`.
|
||||
fn user_id_from_headers(state: &AppState, headers: &HeaderMap) -> Option<i64> {
|
||||
let raw = headers
|
||||
.get(axum::http::header::AUTHORIZATION)
|
||||
.and_then(|v| v.to_str().ok())?;
|
||||
let token = raw.strip_prefix("Bearer ").unwrap_or(raw).trim();
|
||||
user_id_from_token(state, token)
|
||||
}
|
||||
|
||||
fn user_id_from_cookie(state: &AppState, headers: &HeaderMap) -> Option<i64> {
|
||||
let raw = headers
|
||||
.get(axum::http::header::COOKIE)
|
||||
.and_then(|v| v.to_str().ok())?;
|
||||
for part in raw.split(';') {
|
||||
let mut it = part.trim().splitn(2, '=');
|
||||
let name = it.next()?.trim();
|
||||
let value = it.next().unwrap_or("").trim();
|
||||
if name == COMFY_SESSION_COOKIE {
|
||||
let token = urlencoding::decode(value).ok()?;
|
||||
return user_id_from_token(state, token.as_ref());
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
fn token_from_query(raw_query: Option<&str>) -> Option<String> {
|
||||
let q = raw_query?;
|
||||
for pair in q.split('&') {
|
||||
let mut it = pair.splitn(2, '=');
|
||||
let k = it.next()?;
|
||||
let v = it.next().unwrap_or("");
|
||||
if k == "token" {
|
||||
return Some(urlencoding::decode(v).ok()?.into_owned());
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
fn query_without_token(raw_query: Option<&str>) -> Option<String> {
|
||||
let query = raw_query?;
|
||||
let kept: Vec<&str> = query
|
||||
.split('&')
|
||||
.filter(|pair| pair.splitn(2, '=').next() != Some("token"))
|
||||
.filter(|pair| !pair.is_empty())
|
||||
.collect();
|
||||
if kept.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(kept.join("&"))
|
||||
}
|
||||
}
|
||||
|
||||
fn comfy_session_cookie(token: &str) -> String {
|
||||
format!(
|
||||
"{}={}; Path=/comfy; HttpOnly; SameSite=Lax",
|
||||
COMFY_SESSION_COOKIE,
|
||||
urlencoding::encode(token)
|
||||
)
|
||||
}
|
||||
|
||||
// ── HTTP forwarder ─────────────────────────────────────────────────────
|
||||
|
||||
async fn http_handler(State(state): State<AppState>, req: Request<Body>) -> Response<Body> {
|
||||
let query_token = token_from_query(req.uri().query());
|
||||
let user_id = user_id_from_headers(&state, req.headers())
|
||||
.or_else(|| user_id_from_cookie(&state, req.headers()))
|
||||
.or_else(|| {
|
||||
query_token
|
||||
.as_deref()
|
||||
.and_then(|token| user_id_from_token(&state, token))
|
||||
});
|
||||
let Some(user_id) = user_id else {
|
||||
return unauthorized();
|
||||
};
|
||||
|
||||
let entry = match state.comfyui_pool.get_or_start(user_id).await {
|
||||
Ok(e) => e,
|
||||
Err(PoolError::AtCapacity(_)) => return at_capacity(),
|
||||
Err(e) => return upstream_error(format!("pool: {}", e)),
|
||||
};
|
||||
|
||||
if entry.state != ContainerState::Ready {
|
||||
if let Err(e) = wait_for_ready(&state.comfyui_pool, user_id).await {
|
||||
return upstream_error(e);
|
||||
}
|
||||
}
|
||||
|
||||
state.comfyui_pool.touch(user_id).await;
|
||||
|
||||
let path = req.uri().path();
|
||||
let upstream_path = path.strip_prefix("/comfy").unwrap_or(path);
|
||||
let upstream_path = if upstream_path.is_empty() {
|
||||
"/"
|
||||
} else {
|
||||
upstream_path
|
||||
};
|
||||
let query = query_without_token(req.uri().query());
|
||||
|
||||
let port = match state.comfyui_pool.entry(user_id).await {
|
||||
Some(e) => e.host_port,
|
||||
None => return upstream_error("container vanished".into()),
|
||||
};
|
||||
|
||||
let upstream = match build_upstream_uri("http", port, upstream_path, query.as_deref()) {
|
||||
Ok(u) => u,
|
||||
Err(e) => return upstream_error(e),
|
||||
};
|
||||
|
||||
let mut response = forward_http(state.comfyui_pool.clone(), user_id, req, upstream).await;
|
||||
if let Some(token) = query_token {
|
||||
if let Ok(value) = HeaderValue::from_str(&comfy_session_cookie(&token)) {
|
||||
response
|
||||
.headers_mut()
|
||||
.append(axum::http::header::SET_COOKIE, value);
|
||||
}
|
||||
}
|
||||
response
|
||||
}
|
||||
|
||||
fn build_upstream_uri(
|
||||
scheme: &str,
|
||||
port: u16,
|
||||
path: &str,
|
||||
query: Option<&str>,
|
||||
) -> Result<Uri, String> {
|
||||
let mut s = format!("{}://127.0.0.1:{}{}", scheme, port, path);
|
||||
if let Some(q) = query {
|
||||
s.push('?');
|
||||
s.push_str(q);
|
||||
}
|
||||
Uri::try_from(s).map_err(|e| format!("bad upstream uri: {}", e))
|
||||
}
|
||||
|
||||
async fn forward_http(
|
||||
pool: Arc<ComfyUIPool>,
|
||||
user_id: i64,
|
||||
req: Request<Body>,
|
||||
upstream: Uri,
|
||||
) -> Response<Body> {
|
||||
let (parts, body) = req.into_parts();
|
||||
|
||||
let client = match reqwest::Client::builder()
|
||||
.redirect(reqwest::redirect::Policy::none())
|
||||
.timeout(Duration::from_secs(120))
|
||||
.build()
|
||||
{
|
||||
Ok(c) => c,
|
||||
Err(e) => return upstream_error(format!("client: {}", e)),
|
||||
};
|
||||
|
||||
let body_bytes = match axum::body::to_bytes(body, 64 * 1024 * 1024).await {
|
||||
Ok(b) => b,
|
||||
Err(e) => return upstream_error(format!("read body: {}", e)),
|
||||
};
|
||||
|
||||
let mut rb = client.request(reqwest_method(&parts.method), upstream.to_string());
|
||||
for (name, value) in parts.headers.iter() {
|
||||
if is_hop_by_hop(name) || is_browser_context_header(name) {
|
||||
continue;
|
||||
}
|
||||
if let Ok(v) = value.to_str() {
|
||||
rb = rb.header(name.as_str(), v);
|
||||
}
|
||||
}
|
||||
rb = rb.body(body_bytes);
|
||||
|
||||
let resp = match rb.send().await {
|
||||
Ok(r) => r,
|
||||
Err(e) => return upstream_error(format!("upstream send: {}", e)),
|
||||
};
|
||||
|
||||
pool.touch(user_id).await;
|
||||
|
||||
let status = resp.status();
|
||||
let mut builder = Response::builder().status(status);
|
||||
for (name, value) in resp.headers().iter() {
|
||||
if is_hop_by_hop(name) {
|
||||
continue;
|
||||
}
|
||||
if let Ok(v) = HeaderValue::from_bytes(value.as_bytes()) {
|
||||
builder = builder.header(name.as_str(), v);
|
||||
}
|
||||
}
|
||||
|
||||
let bytes = match resp.bytes().await {
|
||||
Ok(b) => b,
|
||||
Err(e) => return upstream_error(format!("upstream read: {}", e)),
|
||||
};
|
||||
|
||||
builder.body(Body::from(bytes)).unwrap_or_else(|_| {
|
||||
Response::builder()
|
||||
.status(StatusCode::BAD_GATEWAY)
|
||||
.body(Body::empty())
|
||||
.unwrap()
|
||||
})
|
||||
}
|
||||
|
||||
fn reqwest_method(m: &Method) -> reqwest::Method {
|
||||
reqwest::Method::from_bytes(m.as_str().as_bytes()).unwrap_or(reqwest::Method::GET)
|
||||
}
|
||||
|
||||
fn is_hop_by_hop(name: &HeaderName) -> bool {
|
||||
name == &CONNECTION
|
||||
|| name == &PROXY_AUTHENTICATE
|
||||
|| name == &PROXY_AUTHORIZATION
|
||||
|| name == &TE
|
||||
|| name == &TRAILER
|
||||
|| name == &TRANSFER_ENCODING
|
||||
|| name == &UPGRADE
|
||||
|| name == &HOST
|
||||
|| name.as_str().eq_ignore_ascii_case("keep-alive")
|
||||
}
|
||||
|
||||
fn is_browser_context_header(name: &HeaderName) -> bool {
|
||||
let name = name.as_str();
|
||||
name.eq_ignore_ascii_case("origin")
|
||||
|| name.eq_ignore_ascii_case("referer")
|
||||
|| name.eq_ignore_ascii_case("sec-fetch-site")
|
||||
|| name.eq_ignore_ascii_case("sec-fetch-mode")
|
||||
|| name.eq_ignore_ascii_case("sec-fetch-dest")
|
||||
|| name.eq_ignore_ascii_case("sec-fetch-user")
|
||||
}
|
||||
|
||||
// ── WebSocket forwarder ────────────────────────────────────────────────
|
||||
|
||||
async fn ws_handler(
|
||||
State(state): State<AppState>,
|
||||
OriginalUri(orig_uri): OriginalUri,
|
||||
RawQuery(raw_query): RawQuery,
|
||||
headers: HeaderMap,
|
||||
ws: WebSocketUpgrade,
|
||||
) -> Response<Body> {
|
||||
// Resolve user from header (preferred) or `?token=` (browser fallback).
|
||||
let user_id = user_id_from_headers(&state, &headers)
|
||||
.or_else(|| user_id_from_cookie(&state, &headers))
|
||||
.or_else(|| {
|
||||
token_from_query(raw_query.as_deref()).and_then(|t| user_id_from_token(&state, &t))
|
||||
});
|
||||
let Some(user_id) = user_id else {
|
||||
return unauthorized();
|
||||
};
|
||||
|
||||
let entry = match state.comfyui_pool.get_or_start(user_id).await {
|
||||
Ok(e) => e,
|
||||
Err(PoolError::AtCapacity(_)) => return at_capacity(),
|
||||
Err(e) => return upstream_error(format!("pool: {}", e)),
|
||||
};
|
||||
if entry.state != ContainerState::Ready {
|
||||
if let Err(e) = wait_for_ready(&state.comfyui_pool, user_id).await {
|
||||
return upstream_error(e);
|
||||
}
|
||||
}
|
||||
let port = match state.comfyui_pool.entry(user_id).await {
|
||||
Some(e) => e.host_port,
|
||||
None => return upstream_error("container vanished".into()),
|
||||
};
|
||||
|
||||
// Forward whatever query the browser sent (clientId etc).
|
||||
let upstream = match build_upstream_uri("ws", port, "/ws", orig_uri.query()) {
|
||||
Ok(u) => u,
|
||||
Err(e) => return upstream_error(e),
|
||||
};
|
||||
|
||||
let pool = state.comfyui_pool.clone();
|
||||
ws.on_upgrade(move |socket| async move {
|
||||
pool.ws_open(user_id).await;
|
||||
if let Err(e) = pump_ws(socket, upstream.to_string(), pool.clone(), user_id).await {
|
||||
warn!(user_id, error = %e, "ws proxy ended with error");
|
||||
}
|
||||
pool.ws_close(user_id).await;
|
||||
})
|
||||
.into_response()
|
||||
}
|
||||
|
||||
async fn pump_ws(
|
||||
client: WebSocket,
|
||||
upstream_url: String,
|
||||
pool: Arc<ComfyUIPool>,
|
||||
user_id: i64,
|
||||
) -> Result<(), String> {
|
||||
let (upstream_ws, _) = tokio_tungstenite::connect_async(&upstream_url)
|
||||
.await
|
||||
.map_err(|e| format!("ws connect: {}", e))?;
|
||||
let (mut up_tx, mut up_rx) = upstream_ws.split();
|
||||
let (mut cl_tx, mut cl_rx) = client.split();
|
||||
|
||||
let pool_a = pool.clone();
|
||||
let pool_b = pool;
|
||||
|
||||
let to_upstream = async move {
|
||||
while let Some(frame) = cl_rx.next().await {
|
||||
let frame = match frame {
|
||||
Ok(f) => f,
|
||||
Err(e) => return Err::<(), String>(format!("client recv: {}", e)),
|
||||
};
|
||||
pool_a.touch(user_id).await;
|
||||
let out = match frame {
|
||||
Message::Text(t) => TMessage::Text(t.to_string().into()),
|
||||
Message::Binary(b) => TMessage::Binary(b.to_vec()),
|
||||
Message::Ping(p) => TMessage::Ping(p.to_vec()),
|
||||
Message::Pong(p) => TMessage::Pong(p.to_vec()),
|
||||
Message::Close(_) => TMessage::Close(None),
|
||||
};
|
||||
if up_tx.send(out).await.is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
};
|
||||
|
||||
let to_client = async move {
|
||||
while let Some(frame) = up_rx.next().await {
|
||||
let frame = match frame {
|
||||
Ok(f) => f,
|
||||
Err(e) => return Err::<(), String>(format!("upstream recv: {}", e)),
|
||||
};
|
||||
pool_b.touch(user_id).await;
|
||||
let out = match frame {
|
||||
TMessage::Text(t) => Message::Text(t.to_string().into()),
|
||||
TMessage::Binary(b) => Message::Binary(b.into()),
|
||||
TMessage::Ping(p) => Message::Ping(p.into()),
|
||||
TMessage::Pong(p) => Message::Pong(p.into()),
|
||||
TMessage::Close(_) => Message::Close(None),
|
||||
TMessage::Frame(_) => continue,
|
||||
};
|
||||
if cl_tx.send(out).await.is_err() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
};
|
||||
|
||||
tokio::select! {
|
||||
r = to_upstream => r,
|
||||
r = to_client => r,
|
||||
}
|
||||
}
|
||||
|
||||
// ── Cold-start hold ────────────────────────────────────────────────────
|
||||
|
||||
async fn wait_for_ready(pool: &ComfyUIPool, user_id: i64) -> Result<(), String> {
|
||||
let deadline = Instant::now() + COLD_START_HOLD;
|
||||
loop {
|
||||
match pool.state(user_id).await {
|
||||
ContainerState::Ready => return Ok(()),
|
||||
ContainerState::Failed | ContainerState::Stopped => {
|
||||
return Err("container failed to start".into());
|
||||
}
|
||||
ContainerState::Starting => {
|
||||
if Instant::now() >= deadline {
|
||||
return Err("cold start timeout".into());
|
||||
}
|
||||
tokio::time::sleep(COLD_START_POLL_INTERVAL).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ── Error responses ────────────────────────────────────────────────────
|
||||
|
||||
fn unauthorized() -> Response<Body> {
|
||||
Response::builder()
|
||||
.status(StatusCode::UNAUTHORIZED)
|
||||
.header("WWW-Authenticate", "Bearer")
|
||||
.body(Body::from(r#"{"error":"unauthorized"}"#))
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
fn at_capacity() -> Response<Body> {
|
||||
Response::builder()
|
||||
.status(StatusCode::SERVICE_UNAVAILABLE)
|
||||
.header("Retry-After", "60")
|
||||
.header("Content-Type", "application/json")
|
||||
.body(Body::from(
|
||||
r#"{"error":"comfyui pool at capacity","retry_after":60}"#,
|
||||
))
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
fn upstream_error(msg: String) -> Response<Body> {
|
||||
debug!(%msg, "comfy proxy upstream error");
|
||||
Response::builder()
|
||||
.status(StatusCode::BAD_GATEWAY)
|
||||
.header("Content-Type", "application/json")
|
||||
.body(Body::from(format!(
|
||||
r#"{{"error":"comfy proxy: {}"}}"#,
|
||||
msg.replace('"', "\\\"")
|
||||
)))
|
||||
.unwrap()
|
||||
}
|
||||
+152
-8
@@ -63,6 +63,7 @@ pub fn routes() -> Router<AppState> {
|
||||
Router::new()
|
||||
.route("/api/comfyui-agent/status", get(status_handler))
|
||||
.route("/api/comfyui-agent/models", get(models_handler))
|
||||
.route("/api/comfyui-agent/image-models", get(image_models_handler))
|
||||
.route("/api/comfyui-agent/chat", post(chat_handler))
|
||||
}
|
||||
|
||||
@@ -70,6 +71,8 @@ pub fn routes() -> Router<AppState> {
|
||||
struct AgentChatRequest {
|
||||
messages: Vec<AgentMessage>,
|
||||
model: Option<String>,
|
||||
#[serde(rename = "imageModel")]
|
||||
image_model: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
@@ -85,6 +88,10 @@ struct AgentMessage {
|
||||
struct AgentAttachment {
|
||||
name: String,
|
||||
mime_type: Option<String>,
|
||||
#[serde(default, rename = "mediaKind")]
|
||||
media_kind: Option<String>,
|
||||
#[serde(default, rename = "referenceLabel")]
|
||||
reference_label: Option<String>,
|
||||
text: Option<String>,
|
||||
data_url: Option<String>,
|
||||
}
|
||||
@@ -290,6 +297,52 @@ fn is_exact_model_pattern(value: &str) -> bool {
|
||||
&& !normalized.contains('?')
|
||||
}
|
||||
|
||||
fn is_image_model_pattern(value: &str) -> bool {
|
||||
if !is_exact_model_pattern(value) {
|
||||
return false;
|
||||
}
|
||||
|
||||
let normalized = value.trim().to_ascii_lowercase();
|
||||
let image_keywords = [
|
||||
"gpt-image",
|
||||
"dall-e",
|
||||
"dalle",
|
||||
"midjourney",
|
||||
"recraft",
|
||||
"kolors",
|
||||
"cogview",
|
||||
"seedream",
|
||||
"sdxl",
|
||||
"stable-diffusion",
|
||||
"stablediffusion",
|
||||
"pixart",
|
||||
"playground",
|
||||
"firefly",
|
||||
"imagen",
|
||||
"flux",
|
||||
"image",
|
||||
"img",
|
||||
"canvas",
|
||||
"painting",
|
||||
"draw",
|
||||
];
|
||||
if image_keywords
|
||||
.iter()
|
||||
.any(|keyword| normalized.contains(keyword))
|
||||
{
|
||||
return true;
|
||||
}
|
||||
|
||||
let chat_keywords = [
|
||||
"claude", "chatgpt", "gpt-4", "gpt-3", "gpt-5", "o1", "o3", "o4", "gemini", "deepseek",
|
||||
"qwen", "llama", "mistral", "command", "kimi", "doubao", "glm", "hunyuan", "spark",
|
||||
"baichuan", "ernie", "wenxin", "grok",
|
||||
];
|
||||
!chat_keywords
|
||||
.iter()
|
||||
.any(|keyword| normalized.contains(keyword))
|
||||
}
|
||||
|
||||
async fn status_handler(
|
||||
State(state): State<AppState>,
|
||||
headers: HeaderMap,
|
||||
@@ -353,6 +406,47 @@ async fn models_handler(
|
||||
})))
|
||||
}
|
||||
|
||||
async fn image_models_handler(
|
||||
State(state): State<AppState>,
|
||||
headers: HeaderMap,
|
||||
) -> Result<Json<Value>, (axum::http::StatusCode, Json<Value>)> {
|
||||
let _auth = resolve_agent_auth(&headers, &state).await?;
|
||||
|
||||
let rows: Vec<String> = sqlx::query_scalar(
|
||||
"SELECT model_pattern
|
||||
FROM token_routes
|
||||
WHERE COALESCE(enabled, 1) = 1
|
||||
ORDER BY model_pattern COLLATE NOCASE",
|
||||
)
|
||||
.fetch_all(&*state.db)
|
||||
.await
|
||||
.map_err(|err| {
|
||||
(
|
||||
axum::http::StatusCode::INTERNAL_SERVER_ERROR,
|
||||
Json(json!({ "error": err.to_string() })),
|
||||
)
|
||||
})?;
|
||||
|
||||
let mut seen = std::collections::HashSet::new();
|
||||
let models: Vec<String> = rows
|
||||
.into_iter()
|
||||
.map(|value| value.trim().to_string())
|
||||
.filter(|value| is_image_model_pattern(value))
|
||||
.filter(|value| seen.insert(value.clone()))
|
||||
.collect();
|
||||
|
||||
Ok(Json(json!({
|
||||
"success": true,
|
||||
"models": models,
|
||||
"defaultModel": models
|
||||
.iter()
|
||||
.find(|model| model.to_ascii_lowercase().contains("image"))
|
||||
.or_else(|| models.first())
|
||||
.cloned()
|
||||
.unwrap_or_else(|| "gpt-image-2".to_string()),
|
||||
})))
|
||||
}
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
struct ParsedSections {
|
||||
characters: Vec<CharacterSection>,
|
||||
@@ -501,7 +595,12 @@ fn parse_sections(full_text: &str) -> ParsedSections {
|
||||
sections
|
||||
}
|
||||
|
||||
async fn call_image_gen(api_base: &str, api_key: &str, prompt: &str) -> Result<String, String> {
|
||||
async fn call_image_gen(
|
||||
api_base: &str,
|
||||
api_key: &str,
|
||||
image_model: &str,
|
||||
prompt: &str,
|
||||
) -> Result<String, String> {
|
||||
let client = Client::new();
|
||||
let url = format!("{}/images/generations", api_base.trim_end_matches('/'));
|
||||
|
||||
@@ -509,7 +608,7 @@ async fn call_image_gen(api_base: &str, api_key: &str, prompt: &str) -> Result<S
|
||||
.post(&url)
|
||||
.bearer_auth(api_key)
|
||||
.json(&json!({
|
||||
"model": "gpt-image-2",
|
||||
"model": image_model,
|
||||
"prompt": prompt,
|
||||
"n": 1,
|
||||
"size": "1024x1024"
|
||||
@@ -587,6 +686,19 @@ fn attachment_summary(attachment: &AgentAttachment) -> String {
|
||||
.unwrap_or("application/octet-stream");
|
||||
let name = attachment.name.trim();
|
||||
let name = if name.is_empty() { "untitled" } else { name };
|
||||
let reference = attachment
|
||||
.reference_label
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.unwrap_or("@素材");
|
||||
let media_kind = attachment
|
||||
.media_kind
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.unwrap_or("file");
|
||||
let header = format!("\n\n[素材引用: {reference}, 名称: {name}, 类型: {mime}, 分类: {media_kind}]");
|
||||
|
||||
if let Some(text) = attachment
|
||||
.text
|
||||
@@ -594,9 +706,9 @@ fn attachment_summary(attachment: &AgentAttachment) -> String {
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
{
|
||||
format!("\n\n[附件: {name}, 类型: {mime}]\n{text}")
|
||||
format!("{header}\n{text}")
|
||||
} else {
|
||||
format!("\n\n[附件: {name}, 类型: {mime}]")
|
||||
header
|
||||
}
|
||||
}
|
||||
|
||||
@@ -684,6 +796,12 @@ async fn chat_handler(
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned);
|
||||
let requested_image_model = body
|
||||
.image_model
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(ToOwned::to_owned);
|
||||
|
||||
let messages: Vec<Value> = body
|
||||
.messages
|
||||
@@ -699,6 +817,7 @@ async fn chat_handler(
|
||||
"comfyui-agent: resolved upstream",
|
||||
);
|
||||
let agent_model = requested_model.unwrap_or_else(|| state.config.comfyui.agent_model.clone());
|
||||
let image_model = requested_image_model.unwrap_or_else(|| "gpt-image-2".to_string());
|
||||
let (tx, rx) = unbounded_channel::<Result<Event, Infallible>>();
|
||||
|
||||
tokio::spawn(async move {
|
||||
@@ -813,12 +932,35 @@ async fn chat_handler(
|
||||
continue;
|
||||
};
|
||||
|
||||
let content = parsed
|
||||
let delta = parsed
|
||||
.get("choices")
|
||||
.and_then(Value::as_array)
|
||||
.and_then(|choices| choices.first())
|
||||
.and_then(|choice| choice.get("delta"))
|
||||
.and_then(|delta| delta.get("content"))
|
||||
.cloned()
|
||||
.unwrap_or(Value::Null);
|
||||
|
||||
let reasoning = delta
|
||||
.get("reasoning_content")
|
||||
.or_else(|| delta.get("reasoning"))
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default()
|
||||
.to_string();
|
||||
|
||||
if !reasoning.is_empty() {
|
||||
if !send_sse(
|
||||
&tx,
|
||||
json!({
|
||||
"type": "reasoning",
|
||||
"content": reasoning,
|
||||
}),
|
||||
) {
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
let content = delta
|
||||
.get("content")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default()
|
||||
.to_string();
|
||||
@@ -893,7 +1035,7 @@ async fn chat_handler(
|
||||
tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
|
||||
}
|
||||
|
||||
match call_image_gen(&api_base, &api_key, prompt).await {
|
||||
match call_image_gen(&api_base, &api_key, &image_model, prompt).await {
|
||||
Ok(url) if !url.is_empty() => {
|
||||
send_sse(
|
||||
&tx,
|
||||
@@ -927,7 +1069,7 @@ async fn chat_handler(
|
||||
|
||||
tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
|
||||
|
||||
match call_image_gen(&api_base, &api_key, &scene.prompt).await {
|
||||
match call_image_gen(&api_base, &api_key, &image_model, &scene.prompt).await {
|
||||
Ok(url) if !url.is_empty() => {
|
||||
send_sse(
|
||||
&tx,
|
||||
@@ -1191,8 +1333,10 @@ mod tests {
|
||||
messages: vec![AgentMessage {
|
||||
role: "user".to_string(),
|
||||
content: "hello".to_string(),
|
||||
attachments: Vec::new(),
|
||||
}],
|
||||
model: None,
|
||||
image_model: None,
|
||||
}),
|
||||
)
|
||||
.await;
|
||||
|
||||
@@ -0,0 +1,160 @@
|
||||
use crate::{
|
||||
auth::jwt::validate_jwt,
|
||||
db::models::downstream_api_key::DownstreamApiKey,
|
||||
error::{AppError, AppResult},
|
||||
AppState,
|
||||
};
|
||||
use axum::{
|
||||
extract::State,
|
||||
http::{header, HeaderMap},
|
||||
routing::{get, post},
|
||||
Json, Router,
|
||||
};
|
||||
use serde_json::{json, Value};
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct DigitalHumanAuth {
|
||||
outbound_key: String,
|
||||
}
|
||||
|
||||
pub fn routes() -> Router<AppState> {
|
||||
Router::new()
|
||||
.route("/api/digital-human/status", get(status))
|
||||
.route("/api/digital-human/generate", post(generate))
|
||||
}
|
||||
|
||||
fn local_v1_base(state: &AppState) -> String {
|
||||
let host = if state.config.server.host == "0.0.0.0" || state.config.server.host == "::" {
|
||||
"127.0.0.1".to_string()
|
||||
} else {
|
||||
state.config.server.host.clone()
|
||||
};
|
||||
format!("http://{}:{}/v1", host, state.config.server.port)
|
||||
}
|
||||
|
||||
fn extract_bearer_token(headers: &HeaderMap) -> Option<&str> {
|
||||
headers
|
||||
.get(header::AUTHORIZATION)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.and_then(|value| {
|
||||
value
|
||||
.strip_prefix("Bearer ")
|
||||
.or_else(|| value.strip_prefix("bearer "))
|
||||
})
|
||||
.map(str::trim)
|
||||
.filter(|token| !token.is_empty())
|
||||
}
|
||||
|
||||
async fn user_downstream_key(state: &AppState, user_id: i64) -> Result<DownstreamApiKey, String> {
|
||||
sqlx::query_as::<_, DownstreamApiKey>(
|
||||
"SELECT *
|
||||
FROM downstream_api_keys
|
||||
WHERE user_id = ?
|
||||
AND COALESCE(enabled, 1) = 1
|
||||
AND (
|
||||
expires_at IS NULL
|
||||
OR expires_at = ''
|
||||
OR datetime(expires_at) > datetime('now')
|
||||
)
|
||||
ORDER BY updated_at DESC, created_at DESC, id DESC
|
||||
LIMIT 1",
|
||||
)
|
||||
.bind(user_id)
|
||||
.fetch_optional(&*state.db)
|
||||
.await
|
||||
.map_err(|err| err.to_string())?
|
||||
.ok_or_else(|| "当前用户没有可用的 API Key，请先在用户中心创建或启用 API Key".to_string())
|
||||
}
|
||||
|
||||
async fn resolve_auth(headers: &HeaderMap, state: &AppState) -> Result<DigitalHumanAuth, AppError> {
|
||||
let token = extract_bearer_token(headers)
|
||||
.ok_or_else(|| AppError::Unauthorized("Missing Authorization header".into()))?;
|
||||
|
||||
if token == state.config.auth.auth_token || token == state.config.auth.proxy_token {
|
||||
return Ok(DigitalHumanAuth {
|
||||
outbound_key: state.config.auth.proxy_token.clone(),
|
||||
});
|
||||
}
|
||||
|
||||
let claims = validate_jwt(&state.config.auth.jwt_secret, token)
|
||||
.map_err(|_| AppError::Unauthorized("Invalid or expired token".into()))?;
|
||||
if claims.role != "admin" && claims.role != "user" {
|
||||
return Err(AppError::Forbidden("Forbidden".into()));
|
||||
}
|
||||
|
||||
let user_id = claims
|
||||
.sub
|
||||
.parse::<i64>()
|
||||
.map_err(|_| AppError::Unauthorized("Invalid user ID in token".into()))?;
|
||||
match user_downstream_key(state, user_id).await {
|
||||
Ok(key) => Ok(DigitalHumanAuth {
|
||||
outbound_key: key.key,
|
||||
}),
|
||||
Err(_) if claims.role == "admin" => Ok(DigitalHumanAuth {
|
||||
outbound_key: state.config.auth.proxy_token.clone(),
|
||||
}),
|
||||
Err(message) => Err(AppError::Forbidden(message)),
|
||||
}
|
||||
}
|
||||
|
||||
async fn status(State(state): State<AppState>, headers: HeaderMap) -> AppResult<Json<Value>> {
|
||||
let _ = resolve_auth(&headers, &state).await?;
|
||||
Ok(Json(json!({
|
||||
"success": true,
|
||||
"mode": "boosapi-video-route",
|
||||
"capabilities": [
|
||||
"text_to_avatar_video",
|
||||
"image_to_avatar_video",
|
||||
"audio_driven_avatar",
|
||||
"multi_person_prompt",
|
||||
"segmented_long_video"
|
||||
]
|
||||
})))
|
||||
}
|
||||
|
||||
async fn generate(
|
||||
State(state): State<AppState>,
|
||||
headers: HeaderMap,
|
||||
Json(mut body): Json<Value>,
|
||||
) -> AppResult<Json<Value>> {
|
||||
let auth = resolve_auth(&headers, &state).await?;
|
||||
if body.get("model").and_then(Value::as_str).unwrap_or("").trim().is_empty() {
|
||||
return Err(AppError::BadRequest("model is required".into()));
|
||||
}
|
||||
if body
|
||||
.get("prompt")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or("")
|
||||
.trim()
|
||||
.is_empty()
|
||||
{
|
||||
return Err(AppError::BadRequest("prompt is required".into()));
|
||||
}
|
||||
if let Some(object) = body.as_object_mut() {
|
||||
object.insert("task_type".to_string(), json!("digital_human"));
|
||||
}
|
||||
let response = reqwest::Client::builder()
|
||||
.timeout(std::time::Duration::from_secs(300))
|
||||
.build()
|
||||
.unwrap_or_else(|_| reqwest::Client::new())
|
||||
.post(format!("{}/videos", local_v1_base(&state)))
|
||||
.bearer_auth(auth.outbound_key)
|
||||
.json(&body)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|err| AppError::UpstreamError(err.to_string()))?;
|
||||
let status = response.status();
|
||||
let text = response
|
||||
.text()
|
||||
.await
|
||||
.map_err(|err| AppError::UpstreamError(err.to_string()))?;
|
||||
let value = serde_json::from_str::<Value>(&text).unwrap_or_else(|_| json!({ "body": text }));
|
||||
if !status.is_success() {
|
||||
return Err(AppError::UpstreamError(value.to_string()));
|
||||
}
|
||||
Ok(Json(json!({
|
||||
"success": true,
|
||||
"request": body,
|
||||
"response": value
|
||||
})))
|
||||
}
|
||||
+120
-48
@@ -1,22 +1,25 @@
|
||||
use crate::AppState;
|
||||
use axum::{
|
||||
body::Body,
|
||||
http::{Request, StatusCode},
|
||||
response::Response,
|
||||
routing::get,
|
||||
Router,
|
||||
http::{header, HeaderValue, Request, StatusCode},
|
||||
response::{IntoResponse, Response},
|
||||
routing::{any, get},
|
||||
Json, Router,
|
||||
};
|
||||
use include_dir::{include_dir, Dir};
|
||||
use serde_json::{json, Value};
|
||||
use std::convert::Infallible;
|
||||
use tower::service_fn;
|
||||
use tower_http::services::ServeDir;
|
||||
|
||||
static EMBEDDED_WEB: Dir<'_> = include_dir!("$CARGO_MANIFEST_DIR/dist/web");
|
||||
|
||||
mod admin;
|
||||
mod auth;
|
||||
mod checkin;
|
||||
#[cfg(feature = "multi_tenant_comfy")]
|
||||
mod comfy_proxy;
|
||||
mod comfyui;
|
||||
mod comfyui_agent;
|
||||
mod digital_human;
|
||||
mod events;
|
||||
mod health;
|
||||
mod monitor;
|
||||
@@ -31,19 +34,12 @@ mod tasks;
|
||||
mod update_center;
|
||||
mod user;
|
||||
mod user_api_keys;
|
||||
mod workspace;
|
||||
// Two implementations of /api/video-agent/*: legacy singleton (default)
|
||||
// vs per-user pool. Exactly one is mounted, picked at compile time by
|
||||
// the `multi_tenant_comfy` feature flag.
|
||||
#[cfg(feature = "multi_tenant_comfy")]
|
||||
mod video_agent;
|
||||
#[cfg(not(feature = "multi_tenant_comfy"))]
|
||||
#[path = "video_agent_legacy.rs"]
|
||||
mod video_agent;
|
||||
mod workflow_ui;
|
||||
mod workspace;
|
||||
|
||||
pub fn build_router(state: AppState) -> Router {
|
||||
let web_dir = std::env::var("BOOSAPI_WEB_DIR").unwrap_or_else(|_| "dist/web".to_string());
|
||||
|
||||
// Build the API router
|
||||
let api_routes = Router::new()
|
||||
.route("/healthz", get(health::health_check))
|
||||
@@ -51,12 +47,12 @@ pub fn build_router(state: AppState) -> Router {
|
||||
.merge(proxy::routes())
|
||||
.merge(picture_agent::routes())
|
||||
.merge(comfyui_agent::routes())
|
||||
.merge(comfyui::routes());
|
||||
|
||||
// /comfy/* reverse proxy is only mounted when the per-user pool is
|
||||
// active; without it there's no per-user container to forward to.
|
||||
#[cfg(feature = "multi_tenant_comfy")]
|
||||
let api_routes = api_routes.merge(comfy_proxy::routes());
|
||||
.merge(digital_human::routes())
|
||||
.merge(comfyui::routes())
|
||||
.merge(crate::workflow::routes::routes())
|
||||
.route("/comfy", any(removed_comfy_workbench))
|
||||
.route("/comfy/", any(removed_comfy_workbench))
|
||||
.route("/comfy/{*path}", any(removed_comfy_workbench));
|
||||
|
||||
let api_routes = api_routes
|
||||
.merge(auth::routes())
|
||||
@@ -67,6 +63,7 @@ pub fn build_router(state: AppState) -> Router {
|
||||
.merge(organizations::routes())
|
||||
.merge(search::routes())
|
||||
.merge(site_announcements::routes())
|
||||
.merge(workflow_ui::routes())
|
||||
.merge(stats::routes())
|
||||
.merge(tasks::routes())
|
||||
.merge(update_center::routes())
|
||||
@@ -76,7 +73,13 @@ pub fn build_router(state: AppState) -> Router {
|
||||
.merge(video_agent::routes())
|
||||
.with_state(state);
|
||||
|
||||
// Static file serving with SPA fallback
|
||||
let router = Router::new().merge(api_routes);
|
||||
match std::env::var("BOOSAPI_WEB_DIR")
|
||||
.ok()
|
||||
.map(|value| value.trim().to_string())
|
||||
.filter(|value| !value.is_empty())
|
||||
{
|
||||
Some(web_dir) => {
|
||||
let web_dir_owned = web_dir.clone();
|
||||
let static_service = ServeDir::new(&web_dir)
|
||||
.append_index_html_on_directories(true)
|
||||
@@ -84,42 +87,111 @@ pub fn build_router(state: AppState) -> Router {
|
||||
let web_dir = web_dir_owned.clone();
|
||||
async move {
|
||||
let path = req.uri().path();
|
||||
|
||||
// API routes that don't match → return JSON error, not index.html
|
||||
if path.starts_with("/api/") || path.starts_with("/v1/") {
|
||||
return Ok::<_, Infallible>(
|
||||
Response::builder()
|
||||
.status(StatusCode::NOT_FOUND)
|
||||
.header("Content-Type", "application/json; charset=utf-8")
|
||||
.body(Body::from(
|
||||
serde_json::json!({"error": "Not found", "path": path}).to_string(),
|
||||
))
|
||||
.unwrap(),
|
||||
);
|
||||
if is_api_path(path) {
|
||||
return Ok::<_, Infallible>(api_not_found_response(path));
|
||||
}
|
||||
|
||||
// SPA fallback: serve index.html for client-side routing
|
||||
let index_path = format!("{}/index.html", web_dir);
|
||||
match tokio::fs::read_to_string(&index_path).await {
|
||||
Ok(html) => Ok::<_, Infallible>(
|
||||
Response::builder()
|
||||
.status(StatusCode::OK)
|
||||
.header("Content-Type", "text/html; charset=utf-8")
|
||||
.body(Body::from(html))
|
||||
.unwrap(),
|
||||
),
|
||||
Ok(html) => Ok::<_, Infallible>(html_response(html)),
|
||||
Err(_) => Ok::<_, Infallible>(
|
||||
Response::builder()
|
||||
.status(StatusCode::NOT_FOUND)
|
||||
.header("Content-Type", "text/plain; charset=utf-8")
|
||||
.header(header::CONTENT_TYPE, "text/plain; charset=utf-8")
|
||||
.body(Body::from(format!("404: {} not found", path)))
|
||||
.unwrap(),
|
||||
.unwrap_or_else(|_| StatusCode::NOT_FOUND.into_response()),
|
||||
),
|
||||
}
|
||||
}
|
||||
}));
|
||||
|
||||
Router::new()
|
||||
.merge(api_routes)
|
||||
.fallback_service(static_service)
|
||||
router.fallback_service(static_service)
|
||||
}
|
||||
None => router.fallback(embedded_static_handler),
|
||||
}
|
||||
}
|
||||
|
||||
fn api_not_found_response(path: &str) -> Response {
|
||||
Response::builder()
|
||||
.status(StatusCode::NOT_FOUND)
|
||||
.header(header::CONTENT_TYPE, "application/json; charset=utf-8")
|
||||
.body(Body::from(
|
||||
serde_json::json!({"error": "Not found", "path": path}).to_string(),
|
||||
))
|
||||
.unwrap_or_else(|_| StatusCode::NOT_FOUND.into_response())
|
||||
}
|
||||
|
||||
async fn embedded_static_handler(req: Request<Body>) -> Response {
|
||||
let path = req.uri().path();
|
||||
if is_api_path(path) {
|
||||
return api_not_found_response(path);
|
||||
}
|
||||
|
||||
let asset_path = embedded_asset_path(path);
|
||||
if let Some(file) = EMBEDDED_WEB.get_file(&asset_path) {
|
||||
return embedded_file_response(&asset_path, file.contents());
|
||||
}
|
||||
|
||||
if let Some(file) = EMBEDDED_WEB.get_file("index.html") {
|
||||
return embedded_file_response("index.html", file.contents());
|
||||
}
|
||||
|
||||
Response::builder()
|
||||
.status(StatusCode::NOT_FOUND)
|
||||
.header(header::CONTENT_TYPE, "text/plain; charset=utf-8")
|
||||
.body(Body::from("embedded dist/web/index.html not found"))
|
||||
.unwrap_or_else(|_| StatusCode::NOT_FOUND.into_response())
|
||||
}
|
||||
|
||||
fn is_api_path(path: &str) -> bool {
|
||||
path.starts_with("/api/") || path.starts_with("/v1/")
|
||||
}
|
||||
|
||||
fn embedded_asset_path(path: &str) -> String {
|
||||
let trimmed = path.trim_start_matches('/');
|
||||
if trimmed.is_empty() || trimmed.ends_with('/') {
|
||||
format!("{trimmed}index.html")
|
||||
} else {
|
||||
trimmed.to_string()
|
||||
}
|
||||
}
|
||||
|
||||
fn html_response(html: String) -> Response {
|
||||
Response::builder()
|
||||
.status(StatusCode::OK)
|
||||
.header(header::CONTENT_TYPE, "text/html; charset=utf-8")
|
||||
.body(Body::from(html))
|
||||
.unwrap_or_else(|_| StatusCode::INTERNAL_SERVER_ERROR.into_response())
|
||||
}
|
||||
|
||||
fn embedded_file_response(path: &str, bytes: &[u8]) -> Response {
|
||||
let content_type = mime_guess::from_path(path)
|
||||
.first_or_octet_stream()
|
||||
.essence_str()
|
||||
.to_string();
|
||||
let mut response = Response::builder()
|
||||
.status(StatusCode::OK)
|
||||
.header(header::CONTENT_TYPE, content_type)
|
||||
.body(Body::from(bytes.to_vec()))
|
||||
.unwrap_or_else(|_| StatusCode::INTERNAL_SERVER_ERROR.into_response());
|
||||
|
||||
if path.starts_with("assets/") || path.contains(".") && !path.ends_with("index.html") {
|
||||
response.headers_mut().insert(
|
||||
header::CACHE_CONTROL,
|
||||
HeaderValue::from_static("public, max-age=31536000, immutable"),
|
||||
);
|
||||
}
|
||||
response
|
||||
}
|
||||
|
||||
async fn removed_comfy_workbench() -> (StatusCode, Json<Value>) {
|
||||
(
|
||||
StatusCode::GONE,
|
||||
Json(json!({
|
||||
"error": {
|
||||
"message": "内置 Docker ComfyUI 工作台已移除。请配置外部 comfyui.base_url 后使用 /api/comfyui/proxy，或打开 /workflows 使用 ComfyUI 风格工作流画布。",
|
||||
"type": "ComfyUIWorkbenchRemoved"
|
||||
},
|
||||
"redirect": "/workflows"
|
||||
})),
|
||||
)
|
||||
}
|
||||
|
||||
+584
-29
@@ -1,12 +1,13 @@
|
||||
use crate::{
|
||||
auth::jwt::validate_jwt,
|
||||
db::models::background_task::BackgroundTask,
|
||||
db::models::downstream_api_key::DownstreamApiKey,
|
||||
error::{AppError, AppResult},
|
||||
AppState,
|
||||
};
|
||||
use axum::{
|
||||
body::Body,
|
||||
extract::{Multipart, Query, State},
|
||||
extract::{Multipart, Path, Query, State},
|
||||
http::{header, HeaderMap, Response, StatusCode},
|
||||
response::IntoResponse,
|
||||
routing::{get, post},
|
||||
@@ -19,6 +20,11 @@ use std::time::Duration;
|
||||
|
||||
const TEMPLATE_GALLERY_URL: &str =
|
||||
"https://weiyyin.github.io/linux-do-gpt-image/linux_do_2036927_prompts_images.filtered.json";
|
||||
const TASK_TYPE_PICTURE_GENERATION: &str = "picture-agent-generation";
|
||||
const TEMPLATE_GALLERY_TIMEOUT_SECS: u64 = 30;
|
||||
const PICTURE_GENERATION_TIMEOUT_SECS: u64 = 300;
|
||||
const PICTURE_GENERATION_RETRY_SECS: u64 = 180;
|
||||
const STALE_RUNNING_JOB_MINUTES: i64 = 12;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct PictureAgentAuth {
|
||||
@@ -33,6 +39,13 @@ pub fn routes() -> Router<AppState> {
|
||||
.route("/api/picture-agent/status", get(status_handler))
|
||||
.route("/api/picture-agent/models", get(models_handler))
|
||||
.route("/api/picture-agent/templates", get(templates_handler))
|
||||
.route(
|
||||
"/api/picture-agent/jobs",
|
||||
get(list_jobs).post(create_generation_job),
|
||||
)
|
||||
.route("/api/picture-agent/jobs/clear", post(clear_jobs))
|
||||
.route("/api/picture-agent/jobs/{id}", get(get_job))
|
||||
.route("/api/picture-agent/jobs/{id}/cancel", post(cancel_job))
|
||||
.route(
|
||||
"/api/picture-agent/images/generations",
|
||||
post(images_generations),
|
||||
@@ -54,6 +67,11 @@ struct GalleryItem {
|
||||
image_host_url: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct JobListQuery {
|
||||
limit: Option<i64>,
|
||||
}
|
||||
|
||||
fn local_v1_base(state: &AppState) -> String {
|
||||
let host = if state.config.server.host == "0.0.0.0" || state.config.server.host == "::" {
|
||||
"127.0.0.1".to_string()
|
||||
@@ -149,6 +167,52 @@ fn is_exact_model_pattern(value: &str) -> bool {
|
||||
&& !normalized.contains('?')
|
||||
}
|
||||
|
||||
fn is_picture_model_pattern(value: &str) -> bool {
|
||||
if !is_exact_model_pattern(value) {
|
||||
return false;
|
||||
}
|
||||
|
||||
let normalized = value.trim().to_ascii_lowercase();
|
||||
let image_keywords = [
|
||||
"gpt-image",
|
||||
"dall-e",
|
||||
"dalle",
|
||||
"midjourney",
|
||||
"recraft",
|
||||
"kolors",
|
||||
"cogview",
|
||||
"seedream",
|
||||
"sdxl",
|
||||
"stable-diffusion",
|
||||
"stablediffusion",
|
||||
"pixart",
|
||||
"playground",
|
||||
"firefly",
|
||||
"imagen",
|
||||
"flux",
|
||||
"image",
|
||||
"img",
|
||||
"canvas",
|
||||
"painting",
|
||||
"draw",
|
||||
];
|
||||
if image_keywords
|
||||
.iter()
|
||||
.any(|keyword| normalized.contains(keyword))
|
||||
{
|
||||
return true;
|
||||
}
|
||||
|
||||
let chat_keywords = [
|
||||
"claude", "chatgpt", "gpt-4", "gpt-3", "gpt-5", "o1", "o3", "o4", "gemini", "deepseek",
|
||||
"qwen", "llama", "mistral", "command", "kimi", "doubao", "glm", "hunyuan", "spark",
|
||||
"baichuan", "ernie", "wenxin", "grok",
|
||||
];
|
||||
!chat_keywords
|
||||
.iter()
|
||||
.any(|keyword| normalized.contains(keyword))
|
||||
}
|
||||
|
||||
async fn status_handler(
|
||||
State(state): State<AppState>,
|
||||
headers: HeaderMap,
|
||||
@@ -183,7 +247,7 @@ async fn models_handler(
|
||||
let models: Vec<String> = rows
|
||||
.into_iter()
|
||||
.map(|value| value.trim().to_string())
|
||||
.filter(|value| is_exact_model_pattern(value))
|
||||
.filter(|value| is_picture_model_pattern(value))
|
||||
.filter(|value| seen.insert(value.clone()))
|
||||
.collect();
|
||||
|
||||
@@ -199,6 +263,30 @@ async fn models_handler(
|
||||
})))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn picture_model_filter_keeps_image_models() {
|
||||
assert!(is_picture_model_pattern("gpt-image-2"));
|
||||
assert!(is_picture_model_pattern("dall-e-3"));
|
||||
assert!(is_picture_model_pattern("flux-schnell"));
|
||||
assert!(is_picture_model_pattern("replicate/recraft-v3"));
|
||||
assert!(is_picture_model_pattern("google/imagen-3"));
|
||||
assert!(is_picture_model_pattern("midjourney"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn picture_model_filter_rejects_chat_models() {
|
||||
assert!(!is_picture_model_pattern("gpt-4o"));
|
||||
assert!(!is_picture_model_pattern("claude-3.7-sonnet"));
|
||||
assert!(!is_picture_model_pattern("gemini-2.5-pro"));
|
||||
assert!(!is_picture_model_pattern("deepseek-chat"));
|
||||
assert!(!is_picture_model_pattern("qwen-plus"));
|
||||
}
|
||||
}
|
||||
|
||||
fn prompt_title(prompt: &str, post_number: Option<i64>) -> String {
|
||||
let clean = prompt
|
||||
.split_whitespace()
|
||||
@@ -255,8 +343,8 @@ async fn templates_handler(
|
||||
Query(query): Query<TemplateQuery>,
|
||||
) -> AppResult<Json<Value>> {
|
||||
let _auth = resolve_picture_agent_auth(&headers, &state).await?;
|
||||
let limit = query.limit.unwrap_or(48).clamp(1, 96);
|
||||
let items = proxy_client()
|
||||
let limit = query.limit.unwrap_or(144).clamp(1, 240);
|
||||
let items = template_client()
|
||||
.get(TEMPLATE_GALLERY_URL)
|
||||
.send()
|
||||
.await
|
||||
@@ -311,9 +399,20 @@ async fn templates_handler(
|
||||
})))
|
||||
}
|
||||
|
||||
fn proxy_client() -> Client {
|
||||
fn template_client() -> Client {
|
||||
Client::builder()
|
||||
.timeout(Duration::from_secs(180))
|
||||
.timeout(Duration::from_secs(TEMPLATE_GALLERY_TIMEOUT_SECS))
|
||||
.build()
|
||||
.unwrap_or_else(|_| Client::new())
|
||||
}
|
||||
|
||||
fn picture_generation_client() -> Client {
|
||||
picture_generation_client_with_timeout(PICTURE_GENERATION_TIMEOUT_SECS)
|
||||
}
|
||||
|
||||
fn picture_generation_client_with_timeout(timeout_secs: u64) -> Client {
|
||||
Client::builder()
|
||||
.timeout(Duration::from_secs(timeout_secs))
|
||||
.build()
|
||||
.unwrap_or_else(|_| Client::new())
|
||||
}
|
||||
@@ -336,6 +435,473 @@ async fn archive_picture_agent_response(
|
||||
.await;
|
||||
}
|
||||
|
||||
async fn execute_generation_payload(
|
||||
state: &AppState,
|
||||
auth: &PictureAgentAuth,
|
||||
body: Value,
|
||||
) -> AppResult<(Value, Option<String>, Option<String>)> {
|
||||
execute_generation_payload_with_timeout(state, auth, body, PICTURE_GENERATION_TIMEOUT_SECS)
|
||||
.await
|
||||
}
|
||||
|
||||
fn classify_upstream_failure(status: Option<StatusCode>, text: &str) -> String {
|
||||
let status_text = status
|
||||
.map(|s| format!("{s}"))
|
||||
.unwrap_or_else(|| "unknown".to_string());
|
||||
let lower = text.to_ascii_lowercase();
|
||||
if lower.contains("524") || lower.contains("timeout") {
|
||||
"图片上游超时，请稍后重试或缩短提示词".to_string()
|
||||
} else if lower.contains("502") || lower.contains("bad gateway") {
|
||||
"图片上游网关异常，请稍后重试".to_string()
|
||||
} else if lower.contains("503") || lower.contains("service unavailable") {
|
||||
"图片上游暂时不可用，请稍后重试".to_string()
|
||||
} else if status.is_some_and(|s| s == StatusCode::GATEWAY_TIMEOUT) {
|
||||
"图片上游请求超时，请稍后重试".to_string()
|
||||
} else {
|
||||
format!("Upstream returned {status_text}: {text}")
|
||||
}
|
||||
}
|
||||
|
||||
async fn execute_generation_payload_with_timeout(
|
||||
state: &AppState,
|
||||
auth: &PictureAgentAuth,
|
||||
body: Value,
|
||||
timeout_secs: u64,
|
||||
) -> AppResult<(Value, Option<String>, Option<String>)> {
|
||||
let prompt = body
|
||||
.get("prompt")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(str::to_string);
|
||||
let model = body
|
||||
.get("model")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(str::to_string);
|
||||
let requested_model = model
|
||||
.clone()
|
||||
.ok_or_else(|| AppError::BadRequest("model is required".into()))?;
|
||||
let response = picture_generation_client_with_timeout(timeout_secs)
|
||||
.post(format!("{}/images/generations", local_v1_base(state)))
|
||||
.bearer_auth(auth.outbound_key.clone())
|
||||
.json(&body)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|err| {
|
||||
let hint = if err.is_timeout() {
|
||||
format!(
|
||||
"图片上游在 {} 分钟内没有返回，请稍后重试、缩短提示词或切换图片模型通道",
|
||||
PICTURE_GENERATION_TIMEOUT_SECS / 60
|
||||
)
|
||||
} else if err.is_connect() {
|
||||
"无法连接本地 BoosAPI 图片路由，请确认后端服务仍在运行".to_string()
|
||||
} else {
|
||||
err.to_string()
|
||||
};
|
||||
AppError::UpstreamError(hint)
|
||||
})?;
|
||||
let status = response.status();
|
||||
let text = response
|
||||
.text()
|
||||
.await
|
||||
.map_err(|err| AppError::UpstreamError(err.to_string()))?;
|
||||
if !status.is_success() {
|
||||
return Err(AppError::UpstreamError(classify_upstream_failure(
|
||||
Some(status),
|
||||
&text,
|
||||
)));
|
||||
}
|
||||
let payload: Value = serde_json::from_str(&text).map_err(|err| {
|
||||
AppError::UpstreamError(format!(
|
||||
"Invalid JSON from upstream: {err}; body: {}",
|
||||
text.chars().take(500).collect::<String>()
|
||||
))
|
||||
})?;
|
||||
Ok((payload, prompt, Some(requested_model)))
|
||||
}
|
||||
|
||||
async fn mark_stale_picture_jobs_failed(state: &AppState) -> AppResult<()> {
|
||||
sqlx::query(
|
||||
"UPDATE background_tasks
|
||||
SET status = 'failed',
|
||||
message = '生图任务已中断',
|
||||
error = '后台任务长时间没有进展，可能是后端重启、上游超时或图片通道无响应。请重新提交。',
|
||||
updated_at = datetime('now'),
|
||||
finished_at = datetime('now')
|
||||
WHERE task_type = ?
|
||||
AND status IN ('pending', 'running')
|
||||
AND datetime(COALESCE(updated_at, created_at)) < datetime('now', ?)",
|
||||
)
|
||||
.bind(TASK_TYPE_PICTURE_GENERATION)
|
||||
.bind(format!("-{STALE_RUNNING_JOB_MINUTES} minutes"))
|
||||
.execute(&*state.db)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn job_visible_to_auth(task: &BackgroundTask, auth: &PictureAgentAuth) -> bool {
|
||||
if auth.role == "admin" && auth.user_id.is_none() {
|
||||
return true;
|
||||
}
|
||||
let Some(raw) = task.result_json.as_deref() else {
|
||||
return false;
|
||||
};
|
||||
let Ok(value) = serde_json::from_str::<Value>(raw) else {
|
||||
return false;
|
||||
};
|
||||
value.get("userId").and_then(Value::as_i64) == auth.user_id
|
||||
}
|
||||
|
||||
fn task_to_json(task: BackgroundTask) -> Value {
|
||||
let result = task
|
||||
.result_json
|
||||
.as_deref()
|
||||
.and_then(|raw| serde_json::from_str::<Value>(raw).ok())
|
||||
.unwrap_or_else(|| json!({}));
|
||||
json!({
|
||||
"id": task.id,
|
||||
"type": task.task_type,
|
||||
"title": task.title,
|
||||
"status": task.status,
|
||||
"message": task.message,
|
||||
"error": task.error,
|
||||
"progress": task.progress.unwrap_or(0),
|
||||
"progressTotal": task.progress_total.unwrap_or(4),
|
||||
"result": result,
|
||||
"createdAt": task.created_at,
|
||||
"updatedAt": task.updated_at,
|
||||
"startedAt": task.started_at,
|
||||
"finishedAt": task.finished_at,
|
||||
})
|
||||
}
|
||||
|
||||
fn is_terminal_job_status(status: &str) -> bool {
|
||||
matches!(status, "succeeded" | "failed" | "cancelled")
|
||||
}
|
||||
|
||||
async fn load_picture_job(state: &AppState, id: &str, auth: &PictureAgentAuth) -> AppResult<Value> {
|
||||
mark_stale_picture_jobs_failed(state).await?;
|
||||
let task: BackgroundTask =
|
||||
sqlx::query_as("SELECT * FROM background_tasks WHERE id = ? AND task_type = ? LIMIT 1")
|
||||
.bind(id)
|
||||
.bind(TASK_TYPE_PICTURE_GENERATION)
|
||||
.fetch_optional(&*state.db)
|
||||
.await?
|
||||
.ok_or_else(|| AppError::NotFound("Picture job not found".into()))?;
|
||||
if !job_visible_to_auth(&task, auth) {
|
||||
return Err(AppError::NotFound("Picture job not found".into()));
|
||||
}
|
||||
Ok(task_to_json(task))
|
||||
}
|
||||
|
||||
async fn update_job_progress(
|
||||
state: &AppState,
|
||||
id: &str,
|
||||
progress: i64,
|
||||
message: &str,
|
||||
) -> AppResult<()> {
|
||||
sqlx::query(
|
||||
"UPDATE background_tasks
|
||||
SET progress = ?, message = ?, updated_at = datetime('now')
|
||||
WHERE id = ? AND status IN ('pending', 'running')",
|
||||
)
|
||||
.bind(progress)
|
||||
.bind(message)
|
||||
.bind(id)
|
||||
.execute(&*state.db)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn is_job_cancelled(state: &AppState, id: &str) -> AppResult<bool> {
|
||||
let status: Option<String> =
|
||||
sqlx::query_scalar("SELECT status FROM background_tasks WHERE id = ? LIMIT 1")
|
||||
.bind(id)
|
||||
.fetch_optional(&*state.db)
|
||||
.await?;
|
||||
Ok(status.as_deref() == Some("cancelled"))
|
||||
}
|
||||
|
||||
async fn create_generation_job(
|
||||
State(state): State<AppState>,
|
||||
headers: HeaderMap,
|
||||
Json(body): Json<Value>,
|
||||
) -> AppResult<Json<Value>> {
|
||||
let auth = resolve_picture_agent_auth(&headers, &state).await?;
|
||||
let prompt = body
|
||||
.get("prompt")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or("")
|
||||
.trim()
|
||||
.to_string();
|
||||
let model = body
|
||||
.get("model")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or("gpt-image-2")
|
||||
.trim()
|
||||
.to_string();
|
||||
if prompt.is_empty() {
|
||||
return Err(AppError::BadRequest("prompt is required".into()));
|
||||
}
|
||||
if model.is_empty() {
|
||||
return Err(AppError::BadRequest("model is required".into()));
|
||||
}
|
||||
|
||||
let task_id = format!("pic_{}", uuid::Uuid::new_v4().simple());
|
||||
let title = prompt.chars().take(30).collect::<String>();
|
||||
let initial_result = json!({
|
||||
"userId": auth.user_id,
|
||||
"role": auth.role,
|
||||
"keyName": auth.key_name,
|
||||
"prompt": prompt,
|
||||
"model": model,
|
||||
"mode": "generate",
|
||||
"request": body,
|
||||
});
|
||||
sqlx::query(
|
||||
"INSERT INTO background_tasks
|
||||
(id, task_type, title, status, message, dedupe_key, result_json, progress, progress_total, started_at)
|
||||
VALUES (?, ?, ?, 'running', '任务已提交，等待路由', ?, ?, 1, 4, datetime('now'))",
|
||||
)
|
||||
.bind(&task_id)
|
||||
.bind(TASK_TYPE_PICTURE_GENERATION)
|
||||
.bind(if title.trim().is_empty() { "生图任务" } else { &title })
|
||||
.bind(format!("picture-agent:{}", auth.user_id.map(|id| id.to_string()).unwrap_or_else(|| "admin".into())))
|
||||
.bind(initial_result.to_string())
|
||||
.execute(&*state.db)
|
||||
.await?;
|
||||
|
||||
let state_for_task = state.clone();
|
||||
let task_id_for_spawn = task_id.clone();
|
||||
let auth_for_task = auth.clone();
|
||||
tokio::spawn(async move {
|
||||
let result = async {
|
||||
update_job_progress(
|
||||
&state_for_task,
|
||||
&task_id_for_spawn,
|
||||
2,
|
||||
"本地图片路由已确认，正在请求上游",
|
||||
)
|
||||
.await?;
|
||||
let generation_result =
|
||||
execute_generation_payload(&state_for_task, &auth_for_task, body.clone()).await;
|
||||
let (payload, prompt, model) = match generation_result {
|
||||
Ok(value) => value,
|
||||
Err(AppError::UpstreamError(message))
|
||||
if message.contains("超时") || message.contains("timeout") =>
|
||||
{
|
||||
update_job_progress(
|
||||
&state_for_task,
|
||||
&task_id_for_spawn,
|
||||
2,
|
||||
"上游超时，正在自动重试一次",
|
||||
)
|
||||
.await?;
|
||||
execute_generation_payload_with_timeout(
|
||||
&state_for_task,
|
||||
&auth_for_task,
|
||||
body,
|
||||
PICTURE_GENERATION_RETRY_SECS,
|
||||
)
|
||||
.await?
|
||||
}
|
||||
Err(err) => return Err(err),
|
||||
};
|
||||
if is_job_cancelled(&state_for_task, &task_id_for_spawn).await? {
|
||||
return Ok(());
|
||||
}
|
||||
update_job_progress(
|
||||
&state_for_task,
|
||||
&task_id_for_spawn,
|
||||
3,
|
||||
"图片已返回，正在归档到工作空间",
|
||||
)
|
||||
.await?;
|
||||
archive_picture_agent_response(
|
||||
&state_for_task,
|
||||
&auth_for_task,
|
||||
prompt.clone(),
|
||||
model.clone(),
|
||||
&payload,
|
||||
)
|
||||
.await;
|
||||
let urls = payload
|
||||
.get("data")
|
||||
.and_then(Value::as_array)
|
||||
.cloned()
|
||||
.unwrap_or_default();
|
||||
let final_result = json!({
|
||||
"userId": auth_for_task.user_id,
|
||||
"role": auth_for_task.role,
|
||||
"keyName": auth_for_task.key_name,
|
||||
"prompt": prompt,
|
||||
"model": model,
|
||||
"mode": "generate",
|
||||
"response": payload,
|
||||
"images": urls,
|
||||
});
|
||||
sqlx::query(
|
||||
"UPDATE background_tasks
|
||||
SET status = 'succeeded', message = '生图完成', result_json = ?, progress = 4,
|
||||
updated_at = datetime('now'), finished_at = datetime('now')
|
||||
WHERE id = ? AND status IN ('pending', 'running')",
|
||||
)
|
||||
.bind(final_result.to_string())
|
||||
.bind(&task_id_for_spawn)
|
||||
.execute(&*state_for_task.db)
|
||||
.await?;
|
||||
Ok::<(), AppError>(())
|
||||
}
|
||||
.await;
|
||||
|
||||
if let Err(err) = result {
|
||||
let _ = sqlx::query(
|
||||
"UPDATE background_tasks
|
||||
SET status = 'failed', message = '生图失败', error = ?, updated_at = datetime('now'), finished_at = datetime('now')
|
||||
WHERE id = ? AND status IN ('pending', 'running')",
|
||||
)
|
||||
.bind(err.to_string())
|
||||
.bind(&task_id_for_spawn)
|
||||
.execute(&*state_for_task.db)
|
||||
.await;
|
||||
}
|
||||
});
|
||||
|
||||
let task = load_picture_job(&state, &task_id, &auth).await?;
|
||||
Ok(Json(json!({ "success": true, "job": task })))
|
||||
}
|
||||
|
||||
async fn cancel_job(
|
||||
State(state): State<AppState>,
|
||||
headers: HeaderMap,
|
||||
Path(id): Path<String>,
|
||||
) -> AppResult<Json<Value>> {
|
||||
let auth = resolve_picture_agent_auth(&headers, &state).await?;
|
||||
let task: BackgroundTask =
|
||||
sqlx::query_as("SELECT * FROM background_tasks WHERE id = ? AND task_type = ? LIMIT 1")
|
||||
.bind(&id)
|
||||
.bind(TASK_TYPE_PICTURE_GENERATION)
|
||||
.fetch_optional(&*state.db)
|
||||
.await?
|
||||
.ok_or_else(|| AppError::NotFound("Picture job not found".into()))?;
|
||||
if !job_visible_to_auth(&task, &auth) {
|
||||
return Err(AppError::NotFound("Picture job not found".into()));
|
||||
}
|
||||
if task.status != "pending" && task.status != "running" {
|
||||
return Ok(Json(json!({
|
||||
"success": true,
|
||||
"job": task_to_json(task),
|
||||
})));
|
||||
}
|
||||
|
||||
sqlx::query(
|
||||
"UPDATE background_tasks
|
||||
SET status = 'cancelled',
|
||||
message = '已取消',
|
||||
error = NULL,
|
||||
updated_at = datetime('now'),
|
||||
finished_at = datetime('now')
|
||||
WHERE id = ? AND task_type = ? AND status IN ('pending', 'running')",
|
||||
)
|
||||
.bind(&id)
|
||||
.bind(TASK_TYPE_PICTURE_GENERATION)
|
||||
.execute(&*state.db)
|
||||
.await?;
|
||||
|
||||
let job = load_picture_job(&state, &id, &auth).await?;
|
||||
Ok(Json(json!({ "success": true, "job": job })))
|
||||
}
|
||||
|
||||
async fn list_jobs(
|
||||
State(state): State<AppState>,
|
||||
headers: HeaderMap,
|
||||
Query(query): Query<JobListQuery>,
|
||||
) -> AppResult<Json<Value>> {
|
||||
let auth = resolve_picture_agent_auth(&headers, &state).await?;
|
||||
mark_stale_picture_jobs_failed(&state).await?;
|
||||
let limit = query.limit.unwrap_or(20).clamp(1, 100);
|
||||
let tasks: Vec<BackgroundTask> = sqlx::query_as(
|
||||
"SELECT * FROM background_tasks
|
||||
WHERE task_type = ?
|
||||
ORDER BY created_at DESC
|
||||
LIMIT ?",
|
||||
)
|
||||
.bind(TASK_TYPE_PICTURE_GENERATION)
|
||||
.bind(limit * 3)
|
||||
.fetch_all(&*state.db)
|
||||
.await?;
|
||||
let jobs: Vec<Value> = tasks
|
||||
.into_iter()
|
||||
.filter(|task| job_visible_to_auth(task, &auth))
|
||||
.take(limit as usize)
|
||||
.map(task_to_json)
|
||||
.collect();
|
||||
Ok(Json(json!({ "success": true, "jobs": jobs })))
|
||||
}
|
||||
|
||||
async fn clear_jobs(State(state): State<AppState>, headers: HeaderMap) -> AppResult<Json<Value>> {
|
||||
let auth = resolve_picture_agent_auth(&headers, &state).await?;
|
||||
mark_stale_picture_jobs_failed(&state).await?;
|
||||
|
||||
let tasks: Vec<BackgroundTask> = sqlx::query_as(
|
||||
"SELECT * FROM background_tasks
|
||||
WHERE task_type = ?
|
||||
ORDER BY created_at DESC",
|
||||
)
|
||||
.bind(TASK_TYPE_PICTURE_GENERATION)
|
||||
.fetch_all(&*state.db)
|
||||
.await?;
|
||||
|
||||
let mut cleared = 0_i64;
|
||||
for task in tasks {
|
||||
if !job_visible_to_auth(&task, &auth) || !is_terminal_job_status(&task.status) {
|
||||
continue;
|
||||
}
|
||||
let affected = sqlx::query(
|
||||
"DELETE FROM background_tasks
|
||||
WHERE id = ? AND task_type = ?",
|
||||
)
|
||||
.bind(&task.id)
|
||||
.bind(TASK_TYPE_PICTURE_GENERATION)
|
||||
.execute(&*state.db)
|
||||
.await?
|
||||
.rows_affected();
|
||||
cleared += i64::try_from(affected).unwrap_or(0);
|
||||
}
|
||||
|
||||
Ok(Json(json!({
|
||||
"success": true,
|
||||
"cleared": cleared,
|
||||
})))
|
||||
}
|
||||
|
||||
async fn get_job(
|
||||
State(state): State<AppState>,
|
||||
headers: HeaderMap,
|
||||
Path(id): Path<String>,
|
||||
) -> AppResult<Json<Value>> {
|
||||
let auth = resolve_picture_agent_auth(&headers, &state).await?;
|
||||
let job = load_picture_job(&state, &id, &auth).await?;
|
||||
Ok(Json(json!({ "success": true, "job": job })))
|
||||
}
|
||||
|
||||
async fn json_response_and_archive(
|
||||
state: &AppState,
|
||||
auth: &PictureAgentAuth,
|
||||
payload: Value,
|
||||
prompt: Option<String>,
|
||||
model: Option<String>,
|
||||
) -> Response<Body> {
|
||||
archive_picture_agent_response(state, auth, prompt, model, &payload).await;
|
||||
Response::builder()
|
||||
.status(StatusCode::OK)
|
||||
.header(header::CONTENT_TYPE, "application/json; charset=utf-8")
|
||||
.body(Body::from(payload.to_string()))
|
||||
.unwrap_or_else(|_| StatusCode::INTERNAL_SERVER_ERROR.into_response())
|
||||
}
|
||||
|
||||
async fn upstream_json_response_to_axum_and_archive(
|
||||
state: &AppState,
|
||||
auth: &PictureAgentAuth,
|
||||
@@ -375,27 +941,8 @@ async fn images_generations(
|
||||
Json(body): Json<Value>,
|
||||
) -> AppResult<Response<Body>> {
|
||||
let auth = resolve_picture_agent_auth(&headers, &state).await?;
|
||||
let prompt = body
|
||||
.get("prompt")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(str::to_string);
|
||||
let model = body
|
||||
.get("model")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.map(str::to_string);
|
||||
let url = format!("{}/images/generations", local_v1_base(&state));
|
||||
let response = proxy_client()
|
||||
.post(url)
|
||||
.bearer_auth(auth.outbound_key.clone())
|
||||
.json(&body)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|err| AppError::UpstreamError(err.to_string()))?;
|
||||
upstream_json_response_to_axum_and_archive(&state, &auth, response, prompt, model).await
|
||||
let (payload, prompt, model) = execute_generation_payload(&state, &auth, body).await?;
|
||||
Ok(json_response_and_archive(&state, &auth, payload, prompt, model).await)
|
||||
}
|
||||
|
||||
async fn images_edits(
|
||||
@@ -455,12 +1002,20 @@ async fn images_edits(
|
||||
}
|
||||
|
||||
let url = format!("{}/images/edits", local_v1_base(&state));
|
||||
let response = proxy_client()
|
||||
let response = picture_generation_client()
|
||||
.post(url)
|
||||
.bearer_auth(auth.outbound_key.clone())
|
||||
.multipart(form)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|err| AppError::UpstreamError(err.to_string()))?;
|
||||
upstream_json_response_to_axum_and_archive(&state, &auth, response, prompt, model).await
|
||||
let response = upstream_json_response_to_axum_and_archive(
|
||||
&state,
|
||||
&auth,
|
||||
response,
|
||||
prompt.clone(),
|
||||
model.clone(),
|
||||
)
|
||||
.await?;
|
||||
Ok(response)
|
||||
}
|
||||
|
||||
+160
-23
@@ -103,8 +103,10 @@ async fn messages_handler(
|
||||
let model = canonical.model.clone();
|
||||
|
||||
if canonical.stream {
|
||||
let chunks = execute_messages_stream_with_conductor(state, key_id, model, canonical).await?;
|
||||
let sse = Sse::new(stream::iter(chunks).map(|chunk| {
|
||||
let chunks =
|
||||
execute_messages_stream_with_conductor(state, key_id, model, canonical).await?;
|
||||
let sse =
|
||||
Sse::new(stream::iter(chunks).map(|chunk| {
|
||||
Ok::<_, Infallible>(axum::response::sse::Event::default().data(chunk))
|
||||
}));
|
||||
Ok(sse.into_response())
|
||||
@@ -371,10 +373,53 @@ async fn images_generations(
|
||||
.ok_or_else(|| AppError::BadRequest("model is required".into()))?
|
||||
.to_string();
|
||||
|
||||
let route = resolve_route(&state, &model).await?;
|
||||
let result = executors::execute_image_generation(state.clone(), route, body).await?;
|
||||
record_billing(&state, key_id, &model, 0, 0, 0.0).await;
|
||||
Ok(Json(result))
|
||||
let started_at = Instant::now();
|
||||
let selected = TokenRouter::new(state.db.clone())
|
||||
.resolve_selected(&model)
|
||||
.await?;
|
||||
let body = body_with_selected_model(body, &selected);
|
||||
let result = executors::execute_image_generation(
|
||||
state.clone(),
|
||||
selected_channel_to_route(selected.clone()),
|
||||
body,
|
||||
)
|
||||
.await;
|
||||
|
||||
match result {
|
||||
Ok(payload) => {
|
||||
let cost = estimate_request_cost(&selected, 0, 0);
|
||||
record_proxy_log_success(
|
||||
&state,
|
||||
key_id,
|
||||
&selected,
|
||||
&model,
|
||||
Some(elapsed_ms_since(started_at)),
|
||||
0,
|
||||
0,
|
||||
cost,
|
||||
1,
|
||||
false,
|
||||
)
|
||||
.await;
|
||||
record_billing(&state, key_id, &model, 0, 0, cost).await;
|
||||
Ok(Json(payload))
|
||||
}
|
||||
Err(err) => {
|
||||
record_proxy_log_failure(
|
||||
&state,
|
||||
key_id,
|
||||
Some(&selected),
|
||||
&model,
|
||||
upstream_status_from_error(&err),
|
||||
Some(elapsed_ms_since(started_at)),
|
||||
Some(&err.to_string()),
|
||||
1,
|
||||
false,
|
||||
)
|
||||
.await;
|
||||
Err(err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// POST /v1/images/edits
|
||||
@@ -452,21 +497,63 @@ async fn images_edits(
|
||||
}
|
||||
|
||||
let model = model.ok_or_else(|| AppError::BadRequest("model is required".into()))?;
|
||||
let route = resolve_route(&state, &model).await?;
|
||||
let started_at = Instant::now();
|
||||
let selected = TokenRouter::new(state.db.clone())
|
||||
.resolve_selected(&model)
|
||||
.await?;
|
||||
|
||||
// Override the multipart `model` field with the conductor-resolved
|
||||
// upstream model name (route may map "dall-e-3" → upstream-specific id).
|
||||
if route.model != model {
|
||||
if selected.actual_model != model {
|
||||
for field in fields.iter_mut() {
|
||||
if field.name == "model" && field.bytes.is_none() {
|
||||
field.text = Some(route.model.clone());
|
||||
field.text = Some(selected.actual_model.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let result = executors::execute_image_edit(state.clone(), route, fields).await?;
|
||||
record_billing(&state, key_id, &model, 0, 0, 0.0).await;
|
||||
Ok(Json(result))
|
||||
let result = executors::execute_image_edit(
|
||||
state.clone(),
|
||||
selected_channel_to_route(selected.clone()),
|
||||
fields,
|
||||
)
|
||||
.await;
|
||||
|
||||
match result {
|
||||
Ok(payload) => {
|
||||
let cost = estimate_request_cost(&selected, 0, 0);
|
||||
record_proxy_log_success(
|
||||
&state,
|
||||
key_id,
|
||||
&selected,
|
||||
&model,
|
||||
Some(elapsed_ms_since(started_at)),
|
||||
0,
|
||||
0,
|
||||
cost,
|
||||
1,
|
||||
false,
|
||||
)
|
||||
.await;
|
||||
record_billing(&state, key_id, &model, 0, 0, cost).await;
|
||||
Ok(Json(payload))
|
||||
}
|
||||
Err(err) => {
|
||||
record_proxy_log_failure(
|
||||
&state,
|
||||
key_id,
|
||||
Some(&selected),
|
||||
&model,
|
||||
upstream_status_from_error(&err),
|
||||
Some(elapsed_ms_since(started_at)),
|
||||
Some(&err.to_string()),
|
||||
1,
|
||||
false,
|
||||
)
|
||||
.await;
|
||||
Err(err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// GET /v1/models
|
||||
@@ -564,6 +651,23 @@ async fn record_billing(
|
||||
}
|
||||
}
|
||||
|
||||
fn estimate_request_cost(
|
||||
selected: &SelectedChannel,
|
||||
prompt_tokens: i64,
|
||||
completion_tokens: i64,
|
||||
) -> f64 {
|
||||
let unit_cost = selected.account_unit_cost.unwrap_or(0.0).max(0.0);
|
||||
if unit_cost <= 0.0 {
|
||||
return 0.0;
|
||||
}
|
||||
let total_tokens = (prompt_tokens.max(0) + completion_tokens.max(0)) as f64;
|
||||
if total_tokens <= 0.0 {
|
||||
unit_cost
|
||||
} else {
|
||||
total_tokens / 1_000_000.0 * unit_cost
|
||||
}
|
||||
}
|
||||
|
||||
async fn record_proxy_log_success(
|
||||
state: &AppState,
|
||||
key_id: Option<i64>,
|
||||
@@ -699,6 +803,31 @@ fn body_with_selected_model(mut body: Value, selected: &SelectedChannel) -> Valu
|
||||
body
|
||||
}
|
||||
|
||||
fn ensure_stream_usage(mut body: Value) -> Value {
|
||||
if let Some(object) = body.as_object_mut() {
|
||||
match object.get_mut("stream_options") {
|
||||
Some(options) => {
|
||||
if let Some(map) = options.as_object_mut() {
|
||||
map.entry("include_usage".to_string())
|
||||
.or_insert_with(|| Value::Bool(true));
|
||||
} else {
|
||||
object.insert(
|
||||
"stream_options".to_string(),
|
||||
serde_json::json!({ "include_usage": true }),
|
||||
);
|
||||
}
|
||||
}
|
||||
None => {
|
||||
object.insert(
|
||||
"stream_options".to_string(),
|
||||
serde_json::json!({ "include_usage": true }),
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
body
|
||||
}
|
||||
|
||||
fn build_conductor(state: &AppState) -> ProxyConductor {
|
||||
let db = state.db.clone();
|
||||
let select_db = db.clone();
|
||||
@@ -797,6 +926,7 @@ async fn execute_messages_with_conductor(
|
||||
attempts,
|
||||
} => {
|
||||
let (prompt, completion) = extract_usage(&response);
|
||||
let cost = estimate_request_cost(&selected, prompt, completion);
|
||||
record_proxy_log_success(
|
||||
&billing_state,
|
||||
key_id,
|
||||
@@ -805,7 +935,7 @@ async fn execute_messages_with_conductor(
|
||||
Some(elapsed_ms_since(started_at)),
|
||||
prompt,
|
||||
completion,
|
||||
0.0,
|
||||
cost,
|
||||
attempts,
|
||||
false,
|
||||
)
|
||||
@@ -816,7 +946,7 @@ async fn execute_messages_with_conductor(
|
||||
&billing_model,
|
||||
prompt,
|
||||
completion,
|
||||
0.0,
|
||||
cost,
|
||||
)
|
||||
.await;
|
||||
Ok(response)
|
||||
@@ -938,6 +1068,7 @@ async fn execute_messages_stream_with_conductor(
|
||||
}
|
||||
};
|
||||
|
||||
let cost = estimate_request_cost(&selected, 0, 0);
|
||||
record_proxy_log_success(
|
||||
&billing_state,
|
||||
key_id,
|
||||
@@ -946,12 +1077,12 @@ async fn execute_messages_stream_with_conductor(
|
||||
Some(elapsed_ms_since(started_at)),
|
||||
0,
|
||||
0,
|
||||
0.0,
|
||||
cost,
|
||||
attempts,
|
||||
true,
|
||||
)
|
||||
.await;
|
||||
record_billing(&billing_state, key_id, &billing_model, 0, 0, 0.0).await;
|
||||
record_billing(&billing_state, key_id, &billing_model, 0, 0, cost).await;
|
||||
Ok(chunks)
|
||||
}
|
||||
|
||||
@@ -1005,6 +1136,7 @@ async fn execute_chat_completions_with_conductor(
|
||||
attempts,
|
||||
} => {
|
||||
let (prompt, completion) = extract_usage(&response);
|
||||
let cost = estimate_request_cost(&selected, prompt, completion);
|
||||
record_proxy_log_success(
|
||||
&billing_state,
|
||||
key_id,
|
||||
@@ -1013,7 +1145,7 @@ async fn execute_chat_completions_with_conductor(
|
||||
Some(elapsed_ms_since(started_at)),
|
||||
prompt,
|
||||
completion,
|
||||
0.0,
|
||||
cost,
|
||||
attempts,
|
||||
false,
|
||||
)
|
||||
@@ -1024,7 +1156,7 @@ async fn execute_chat_completions_with_conductor(
|
||||
&billing_model,
|
||||
prompt,
|
||||
completion,
|
||||
0.0,
|
||||
cost,
|
||||
)
|
||||
.await;
|
||||
Ok(response)
|
||||
@@ -1073,7 +1205,8 @@ async fn execute_chat_completions_stream_with_conductor(
|
||||
downstream_policy: None,
|
||||
attempt_fn: Box::new(move |ctx| {
|
||||
let state = state.clone();
|
||||
let body = body_with_selected_model(body.clone(), &ctx.selected);
|
||||
let body =
|
||||
ensure_stream_usage(body_with_selected_model(body.clone(), &ctx.selected));
|
||||
async move {
|
||||
match executors::execute_chat_completions_stream(
|
||||
state,
|
||||
@@ -1160,8 +1293,11 @@ async fn execute_chat_completions_stream_with_conductor(
|
||||
.iter()
|
||||
.rev()
|
||||
.find_map(|chunk| {
|
||||
let payload = chunk.strip_prefix("data: ")?;
|
||||
if payload.trim() == "[DONE]" {
|
||||
let payload = chunk
|
||||
.strip_prefix("data: ")
|
||||
.or_else(|| chunk.strip_prefix("data:"))?
|
||||
.trim();
|
||||
if payload == "[DONE]" {
|
||||
return None;
|
||||
}
|
||||
let value: Value = serde_json::from_str(payload).ok()?;
|
||||
@@ -1173,6 +1309,7 @@ async fn execute_chat_completions_stream_with_conductor(
|
||||
}
|
||||
})
|
||||
.unwrap_or((0, 0));
|
||||
let cost = estimate_request_cost(&selected, prompt, completion);
|
||||
record_proxy_log_success(
|
||||
&billing_state,
|
||||
key_id,
|
||||
@@ -1181,7 +1318,7 @@ async fn execute_chat_completions_stream_with_conductor(
|
||||
Some(elapsed_ms_since(started_at)),
|
||||
prompt,
|
||||
completion,
|
||||
0.0,
|
||||
cost,
|
||||
attempts,
|
||||
true,
|
||||
)
|
||||
@@ -1192,7 +1329,7 @@ async fn execute_chat_completions_stream_with_conductor(
|
||||
&billing_model,
|
||||
prompt,
|
||||
completion,
|
||||
0.0,
|
||||
cost,
|
||||
)
|
||||
.await;
|
||||
Ok(chunks)
|
||||
|
||||
@@ -49,13 +49,13 @@ pub fn routes() -> Router<AppState> {
|
||||
put(update_announcement).delete(delete_announcement),
|
||||
)
|
||||
.route("/api/site-announcements/{id}/read", post(mark_read_admin))
|
||||
.route("/api/site-announcements/read-all", post(mark_all_read_admin))
|
||||
.route(
|
||||
"/api/site-announcements/read-all",
|
||||
post(mark_all_read_admin),
|
||||
)
|
||||
.route("/api/user/announcements", get(list_user_announcements))
|
||||
.route("/api/user/announcements/{id}/read", post(mark_read_user))
|
||||
.route(
|
||||
"/api/user/announcements/read-all",
|
||||
post(mark_all_read_user),
|
||||
)
|
||||
.route("/api/user/announcements/read-all", post(mark_all_read_user))
|
||||
}
|
||||
|
||||
async fn list_admin_announcements(
|
||||
@@ -95,11 +95,7 @@ async fn list_admin_announcements(
|
||||
if let Some(status) = status {
|
||||
query = query.bind(status);
|
||||
}
|
||||
let rows = query
|
||||
.bind(limit)
|
||||
.bind(offset)
|
||||
.fetch_all(&*state.db)
|
||||
.await?;
|
||||
let rows = query.bind(limit).bind(offset).fetch_all(&*state.db).await?;
|
||||
|
||||
Ok(Json(Value::Array(
|
||||
rows.into_iter().map(|row| row.into_admin_json()).collect(),
|
||||
@@ -135,7 +131,9 @@ async fn list_user_announcements(
|
||||
sql.push_str(" AND r.read_at IS NULL");
|
||||
}
|
||||
|
||||
sql.push_str(" ORDER BY COALESCE(a.published_at, a.created_at) DESC, a.id DESC LIMIT ? OFFSET ?");
|
||||
sql.push_str(
|
||||
" ORDER BY COALESCE(a.published_at, a.created_at) DESC, a.id DESC LIMIT ? OFFSET ?",
|
||||
);
|
||||
|
||||
let rows: Vec<UserAnnouncementRow> = sqlx::query_as(&sql)
|
||||
.bind(user_id)
|
||||
@@ -412,10 +410,14 @@ fn normalize_payload(payload: AnnouncementPayload) -> AppResult<NormalizedAnnoun
|
||||
let title = payload.title.trim().to_string();
|
||||
let content = payload.content.trim().to_string();
|
||||
if title.is_empty() {
|
||||
return Err(crate::error::AppError::BadRequest("title is required".into()));
|
||||
return Err(crate::error::AppError::BadRequest(
|
||||
"title is required".into(),
|
||||
));
|
||||
}
|
||||
if content.is_empty() {
|
||||
return Err(crate::error::AppError::BadRequest("content is required".into()));
|
||||
return Err(crate::error::AppError::BadRequest(
|
||||
"content is required".into(),
|
||||
));
|
||||
}
|
||||
|
||||
let level = payload
|
||||
|
||||
+369
-91
@@ -1,5 +1,6 @@
|
||||
use crate::{
|
||||
auth::jwt::validate_jwt,
|
||||
db::models::background_task::BackgroundTask,
|
||||
db::models::{downstream_api_key::DownstreamApiKey, user::User},
|
||||
error::{AppError, AppResult},
|
||||
AppState,
|
||||
@@ -12,6 +13,7 @@ use axum::{
|
||||
};
|
||||
use serde::Deserialize;
|
||||
use serde_json::{json, Value};
|
||||
use sqlx::Row;
|
||||
|
||||
pub fn routes() -> Router<AppState> {
|
||||
Router::new()
|
||||
@@ -20,6 +22,305 @@ pub fn routes() -> Router<AppState> {
|
||||
.route("/v1/user/api-keys/{id}", delete(delete_api_key))
|
||||
.route("/api/user/logs", get(get_user_logs))
|
||||
.route("/api/user/summary", get(get_user_summary))
|
||||
.route(
|
||||
"/api/user/models/marketplace",
|
||||
get(get_user_models_marketplace),
|
||||
)
|
||||
}
|
||||
|
||||
async fn user_picture_log_items(state: &AppState, user_id: i64) -> AppResult<Vec<Value>> {
|
||||
let tasks: Vec<BackgroundTask> = sqlx::query_as(
|
||||
"SELECT *
|
||||
FROM background_tasks
|
||||
WHERE task_type = ?
|
||||
AND json_extract(result_json, '$.userId') = ?
|
||||
ORDER BY created_at DESC, id DESC",
|
||||
)
|
||||
.bind("picture-agent-generation")
|
||||
.bind(user_id)
|
||||
.fetch_all(&*state.db)
|
||||
.await?;
|
||||
|
||||
Ok(tasks
|
||||
.into_iter()
|
||||
.map(|task| {
|
||||
let result = task
|
||||
.result_json
|
||||
.as_deref()
|
||||
.and_then(|raw| serde_json::from_str::<Value>(raw).ok())
|
||||
.unwrap_or_else(|| json!({}));
|
||||
let status = match task.status.as_str() {
|
||||
"succeeded" => "success",
|
||||
"failed" => "failed",
|
||||
"cancelled" => "cancelled",
|
||||
"pending" => "pending",
|
||||
_ => "running",
|
||||
};
|
||||
json!({
|
||||
"id": task.id,
|
||||
"route_id": Value::Null,
|
||||
"channel_id": Value::Null,
|
||||
"account_id": Value::Null,
|
||||
"downstream_api_key_id": Value::Null,
|
||||
"site_id": Value::Null,
|
||||
"username": Value::Null,
|
||||
"site_name": Value::Null,
|
||||
"site_url": Value::Null,
|
||||
"downstream_key_name": result.get("keyName").cloned().unwrap_or(Value::Null),
|
||||
"downstream_key_group_name": Value::Null,
|
||||
"downstream_key_tags": Value::Null,
|
||||
"model_requested": result.get("model").cloned().unwrap_or(Value::Null),
|
||||
"model_actual": result.get("model").cloned().unwrap_or(Value::Null),
|
||||
"status": status,
|
||||
"http_status": if status == "success" { 200 } else if status == "running" || status == "pending" { 202 } else { 500 },
|
||||
"latency_ms": Value::Null,
|
||||
"first_byte_latency_ms": Value::Null,
|
||||
"prompt_tokens": Value::Null,
|
||||
"completion_tokens": Value::Null,
|
||||
"total_tokens": Value::Null,
|
||||
"estimated_cost": 0.0,
|
||||
"error_message": task.error.clone(),
|
||||
"retry_count": 0,
|
||||
"is_stream": false,
|
||||
"client_family": "picture-agent",
|
||||
"client_app_id": "picture-agent",
|
||||
"client_app_name": "Picture Agent",
|
||||
"client_confidence": "high",
|
||||
"created_at": task.created_at,
|
||||
"source": "picture-agent",
|
||||
"task_type": task.task_type,
|
||||
"title": task.title,
|
||||
"message": task.message,
|
||||
"error": task.error,
|
||||
"progress": task.progress,
|
||||
"progress_total": task.progress_total,
|
||||
"status_label": status,
|
||||
"result": result,
|
||||
})
|
||||
})
|
||||
.collect())
|
||||
}
|
||||
|
||||
async fn user_usage_log_items(state: &AppState, user_id: i64) -> AppResult<Vec<Value>> {
|
||||
let rows = sqlx::query(
|
||||
r#"
|
||||
SELECT
|
||||
l.id,
|
||||
l.model,
|
||||
l.cost,
|
||||
l.prompt_tokens,
|
||||
l.completion_tokens,
|
||||
l.total_tokens,
|
||||
l.created_at,
|
||||
k.name AS downstream_key_name
|
||||
FROM usage_logs l
|
||||
JOIN downstream_api_keys k ON k.id = l.downstream_key_id
|
||||
WHERE k.user_id = ?
|
||||
ORDER BY l.created_at DESC, l.id DESC
|
||||
LIMIT 5000
|
||||
"#,
|
||||
)
|
||||
.bind(user_id)
|
||||
.fetch_all(&*state.db)
|
||||
.await?;
|
||||
|
||||
Ok(rows
|
||||
.into_iter()
|
||||
.map(|row| {
|
||||
let model = row
|
||||
.try_get::<Option<String>, _>("model")
|
||||
.ok()
|
||||
.flatten()
|
||||
.unwrap_or_default();
|
||||
let model_trimmed = model.trim().to_string();
|
||||
let source_model = model_trimmed
|
||||
.split(':')
|
||||
.next_back()
|
||||
.unwrap_or(&model_trimmed)
|
||||
.trim()
|
||||
.to_string();
|
||||
json!({
|
||||
"id": format!("usage-{}", row.try_get::<i64, _>("id").unwrap_or_default()),
|
||||
"route_id": Value::Null,
|
||||
"channel_id": Value::Null,
|
||||
"account_id": Value::Null,
|
||||
"downstream_api_key_id": Value::Null,
|
||||
"site_id": Value::Null,
|
||||
"username": Value::Null,
|
||||
"site_name": Value::Null,
|
||||
"site_url": Value::Null,
|
||||
"downstream_key_name": row.try_get::<Option<String>, _>("downstream_key_name").ok().flatten(),
|
||||
"downstream_key_group_name": Value::Null,
|
||||
"downstream_key_tags": Value::Null,
|
||||
"model_requested": source_model,
|
||||
"model_actual": source_model,
|
||||
"status": "success",
|
||||
"http_status": 200,
|
||||
"latency_ms": Value::Null,
|
||||
"first_byte_latency_ms": Value::Null,
|
||||
"prompt_tokens": row.try_get::<Option<i64>, _>("prompt_tokens").ok().flatten(),
|
||||
"completion_tokens": row.try_get::<Option<i64>, _>("completion_tokens").ok().flatten(),
|
||||
"total_tokens": row.try_get::<Option<i64>, _>("total_tokens").ok().flatten(),
|
||||
"estimated_cost": row.try_get::<Option<f64>, _>("cost").ok().flatten().unwrap_or(0.0),
|
||||
"error_message": Value::Null,
|
||||
"retry_count": 0,
|
||||
"is_stream": false,
|
||||
"client_family": "picture-agent",
|
||||
"client_app_id": "picture-agent",
|
||||
"client_app_name": "Picture Agent",
|
||||
"client_confidence": "high",
|
||||
"created_at": row.try_get::<Option<String>, _>("created_at").ok().flatten(),
|
||||
"source": "usage-logs",
|
||||
"task_type": "usage-log",
|
||||
"title": "生图助手",
|
||||
"message": model_trimmed,
|
||||
"error": Value::Null,
|
||||
"progress": Value::Null,
|
||||
"progress_total": Value::Null,
|
||||
})
|
||||
})
|
||||
.collect())
|
||||
}
|
||||
|
||||
fn user_log_text(item: &Value) -> String {
|
||||
[
|
||||
item.get("model_requested").and_then(Value::as_str),
|
||||
item.get("model_actual").and_then(Value::as_str),
|
||||
item.get("downstream_key_name").and_then(Value::as_str),
|
||||
item.get("site_name").and_then(Value::as_str),
|
||||
item.get("site_url").and_then(Value::as_str),
|
||||
item.get("error_message").and_then(Value::as_str),
|
||||
item.get("message").and_then(Value::as_str),
|
||||
item.get("title").and_then(Value::as_str),
|
||||
item.get("source").and_then(Value::as_str),
|
||||
]
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.collect::<Vec<_>>()
|
||||
.join(" ")
|
||||
.to_ascii_lowercase()
|
||||
}
|
||||
|
||||
fn user_log_status(item: &Value) -> String {
|
||||
item.get("status")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default()
|
||||
.to_ascii_lowercase()
|
||||
}
|
||||
|
||||
fn user_log_created_at(item: &Value) -> String {
|
||||
item.get("created_at")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default()
|
||||
.to_string()
|
||||
}
|
||||
|
||||
fn user_log_cost(item: &Value) -> f64 {
|
||||
item.get("estimated_cost")
|
||||
.and_then(Value::as_f64)
|
||||
.unwrap_or(0.0)
|
||||
}
|
||||
|
||||
fn user_log_tokens(item: &Value) -> i64 {
|
||||
item.get("total_tokens")
|
||||
.and_then(Value::as_i64)
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
fn matches_user_log_filters(
|
||||
item: &Value,
|
||||
q: Option<&str>,
|
||||
status: Option<&str>,
|
||||
date_from: Option<&str>,
|
||||
date_to: Option<&str>,
|
||||
) -> bool {
|
||||
if let Some(q) = q {
|
||||
let q = q.to_ascii_lowercase();
|
||||
if !user_log_text(item).contains(&q) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(status) = status {
|
||||
let item_status = user_log_status(item);
|
||||
let normalized = status.to_ascii_lowercase();
|
||||
if normalized == "success" {
|
||||
if item_status != "success" {
|
||||
return false;
|
||||
}
|
||||
} else if normalized == "failed" {
|
||||
if !matches!(item_status.as_str(), "failed" | "cancelled") {
|
||||
return false;
|
||||
}
|
||||
} else if item_status != normalized {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
let created_at = user_log_created_at(item);
|
||||
if let Some(date_from) = date_from {
|
||||
if !created_at.is_empty() && created_at.as_str() < date_from {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
if let Some(date_to) = date_to {
|
||||
if !created_at.is_empty() && created_at.as_str() > date_to {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
true
|
||||
}
|
||||
|
||||
fn sort_user_logs_desc(items: &mut [Value]) {
|
||||
items.sort_by(|a, b| {
|
||||
let created_cmp = user_log_created_at(b).cmp(&user_log_created_at(a));
|
||||
if created_cmp != std::cmp::Ordering::Equal {
|
||||
return created_cmp;
|
||||
}
|
||||
let id_a = a.get("id").map(Value::to_string).unwrap_or_default();
|
||||
let id_b = b.get("id").map(Value::to_string).unwrap_or_default();
|
||||
id_b.cmp(&id_a)
|
||||
});
|
||||
}
|
||||
|
||||
async fn load_user_log_items(state: &AppState, user_id: i64) -> AppResult<Vec<Value>> {
|
||||
let mut items: Vec<Value> = sqlx::query_as::<_, crate::routes::stats::ProxyLogRow>(
|
||||
r#"
|
||||
SELECT
|
||||
l.id, l.route_id, l.channel_id, l.account_id, l.downstream_api_key_id,
|
||||
l.model_requested, l.model_actual, l.status, l.http_status,
|
||||
l.latency_ms, l.first_byte_latency_ms,
|
||||
l.prompt_tokens, l.completion_tokens, l.total_tokens, l.estimated_cost,
|
||||
l.error_message, l.retry_count, l.is_stream,
|
||||
l.client_family, l.client_app_id, l.client_app_name, l.client_confidence,
|
||||
l.created_at,
|
||||
a.site_id AS site_id,
|
||||
COALESCE(NULLIF(a.username, ''), CASE WHEN a.id IS NULL THEN NULL ELSE '账号 #' || a.id END) AS username,
|
||||
s.name AS site_name,
|
||||
s.url AS site_url,
|
||||
k.name AS downstream_key_name,
|
||||
k.group_name AS downstream_key_group_name,
|
||||
k.tags AS downstream_key_tags
|
||||
FROM proxy_logs l
|
||||
JOIN downstream_api_keys k ON l.downstream_api_key_id = k.id
|
||||
LEFT JOIN accounts a ON a.id = l.account_id
|
||||
LEFT JOIN sites s ON s.id = a.site_id
|
||||
WHERE k.user_id = ?
|
||||
ORDER BY l.created_at DESC
|
||||
LIMIT 5000
|
||||
"#,
|
||||
)
|
||||
.bind(user_id)
|
||||
.fetch_all(&*state.db)
|
||||
.await?
|
||||
.into_iter()
|
||||
.map(|row| serde_json::to_value(row).unwrap_or(json!({})))
|
||||
.collect();
|
||||
|
||||
items.extend(user_usage_log_items(state, user_id).await?);
|
||||
items.extend(user_picture_log_items(state, user_id).await?);
|
||||
sort_user_logs_desc(&mut items);
|
||||
Ok(items)
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
@@ -190,98 +491,40 @@ async fn get_user_logs(
|
||||
let offset = params.offset.unwrap_or((page - 1) * page_size).max(0);
|
||||
|
||||
let q = params.q.as_deref().map(str::trim).filter(|s| !s.is_empty());
|
||||
let q_like = q.map(|s| format!("%{s}%"));
|
||||
let status = params
|
||||
.status
|
||||
.as_deref()
|
||||
.filter(|s| !s.is_empty() && *s != "all");
|
||||
let date_from = params.date_from.as_deref().filter(|s| !s.is_empty());
|
||||
let date_to = params.date_to.as_deref().filter(|s| !s.is_empty());
|
||||
|
||||
let items: Vec<Value> = sqlx::query_as::<_, crate::routes::stats::ProxyLogRow>(
|
||||
r#"
|
||||
SELECT
|
||||
l.id, l.route_id, l.channel_id, l.account_id, l.downstream_api_key_id,
|
||||
l.model_requested, l.model_actual, l.status, l.http_status,
|
||||
l.latency_ms, l.first_byte_latency_ms,
|
||||
l.prompt_tokens, l.completion_tokens, l.total_tokens, l.estimated_cost,
|
||||
l.error_message, l.retry_count, l.is_stream,
|
||||
l.client_family, l.client_app_id, l.client_app_name, l.client_confidence,
|
||||
l.created_at,
|
||||
a.site_id AS site_id,
|
||||
COALESCE(NULLIF(a.username, ''), CASE WHEN a.id IS NULL THEN NULL ELSE '账号 #' || a.id END) AS username,
|
||||
s.name AS site_name,
|
||||
s.url AS site_url,
|
||||
k.name AS downstream_key_name,
|
||||
k.group_name AS downstream_key_group_name,
|
||||
k.tags AS downstream_key_tags
|
||||
FROM proxy_logs l
|
||||
JOIN downstream_api_keys k ON l.downstream_api_key_id = k.id
|
||||
LEFT JOIN accounts a ON a.id = l.account_id
|
||||
LEFT JOIN sites s ON s.id = a.site_id
|
||||
WHERE k.user_id = ?
|
||||
AND (? IS NULL OR l.model_requested LIKE ?)
|
||||
AND (? IS NULL OR l.status = ?)
|
||||
AND (? IS NULL OR l.created_at >= ?)
|
||||
AND (? IS NULL OR l.created_at <= ?)
|
||||
ORDER BY l.created_at DESC
|
||||
LIMIT ? OFFSET ?
|
||||
"#,
|
||||
)
|
||||
.bind(user_id)
|
||||
.bind(q_like.as_deref())
|
||||
.bind(q_like.as_deref())
|
||||
.bind(status)
|
||||
.bind(status)
|
||||
.bind(date_from)
|
||||
.bind(date_from)
|
||||
.bind(date_to)
|
||||
.bind(date_to)
|
||||
.bind(page_size)
|
||||
.bind(offset)
|
||||
.fetch_all(&*state.db)
|
||||
.await?
|
||||
let items = load_user_log_items(&state, user_id).await?;
|
||||
let filtered_items: Vec<Value> = items
|
||||
.into_iter()
|
||||
.map(|row| serde_json::to_value(row).unwrap_or(json!({})))
|
||||
.filter(|item| matches_user_log_filters(item, q, status, date_from, date_to))
|
||||
.collect();
|
||||
|
||||
let (total, success_count, total_cost, total_tokens): (i64, i64, f64, i64) = sqlx::query_as(
|
||||
r#"
|
||||
SELECT
|
||||
COUNT(*),
|
||||
SUM(CASE WHEN l.status = 'success' THEN 1 ELSE 0 END),
|
||||
COALESCE(SUM(l.estimated_cost), 0.0),
|
||||
COALESCE(SUM(l.total_tokens), 0)
|
||||
FROM proxy_logs l
|
||||
JOIN downstream_api_keys k ON l.downstream_api_key_id = k.id
|
||||
WHERE k.user_id = ?
|
||||
AND (? IS NULL OR l.model_requested LIKE ?)
|
||||
AND (? IS NULL OR l.status = ?)
|
||||
AND (? IS NULL OR l.created_at >= ?)
|
||||
AND (? IS NULL OR l.created_at <= ?)
|
||||
"#,
|
||||
)
|
||||
.bind(user_id)
|
||||
.bind(q_like.as_deref())
|
||||
.bind(q_like.as_deref())
|
||||
.bind(status)
|
||||
.bind(status)
|
||||
.bind(date_from)
|
||||
.bind(date_from)
|
||||
.bind(date_to)
|
||||
.bind(date_to)
|
||||
.fetch_one(&*state.db)
|
||||
.await?;
|
||||
let total = filtered_items.len() as i64;
|
||||
let success_count = filtered_items
|
||||
.iter()
|
||||
.filter(|item| user_log_status(item).as_str() == "success")
|
||||
.count() as i64;
|
||||
let failed_count = total - success_count;
|
||||
let total_cost = filtered_items.iter().map(user_log_cost).sum::<f64>();
|
||||
let total_tokens = filtered_items.iter().map(user_log_tokens).sum::<i64>();
|
||||
let paged_items = filtered_items
|
||||
.into_iter()
|
||||
.skip(offset as usize)
|
||||
.take(page_size as usize)
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
Ok(Json(json!({
|
||||
"items": items,
|
||||
"items": paged_items,
|
||||
"total": total,
|
||||
"page": page,
|
||||
"page_size": page_size,
|
||||
"summary": {
|
||||
"total": total,
|
||||
"success": success_count,
|
||||
"failed": total - success_count,
|
||||
"failed": failed_count,
|
||||
"total_cost": total_cost,
|
||||
"total_tokens": total_tokens,
|
||||
}
|
||||
@@ -293,22 +536,57 @@ async fn get_user_summary(
|
||||
headers: HeaderMap,
|
||||
) -> AppResult<Json<Value>> {
|
||||
let user_id = current_user_id(&headers, &state)?;
|
||||
|
||||
let stats: (i64, i64, f64) = sqlx::query_as(
|
||||
"SELECT COUNT(*),
|
||||
SUM(CASE WHEN l.status = 'success' THEN 1 ELSE 0 END),
|
||||
COALESCE(SUM(l.estimated_cost), 0.0)
|
||||
FROM proxy_logs l
|
||||
JOIN downstream_api_keys k ON l.downstream_api_key_id = k.id
|
||||
WHERE k.user_id = ?",
|
||||
)
|
||||
.bind(user_id)
|
||||
.fetch_one(&*state.db)
|
||||
.await?;
|
||||
let items = load_user_log_items(&state, user_id).await?;
|
||||
let total = items.len() as i64;
|
||||
let success_count = items
|
||||
.iter()
|
||||
.filter(|item| user_log_status(item).as_str() == "success")
|
||||
.count() as i64;
|
||||
let total_cost = items.iter().map(user_log_cost).sum::<f64>();
|
||||
|
||||
Ok(Json(json!({
|
||||
"totalRequests": stats.0,
|
||||
"successRequests": stats.1,
|
||||
"totalCost": stats.2,
|
||||
"totalRequests": total,
|
||||
"successRequests": success_count,
|
||||
"totalCost": total_cost,
|
||||
})))
|
||||
}
|
||||
|
||||
async fn get_user_models_marketplace(
|
||||
State(state): State<AppState>,
|
||||
headers: HeaderMap,
|
||||
Query(query): Query<crate::routes::admin::models_marketplace::MarketplaceQuery>,
|
||||
) -> AppResult<Json<Value>> {
|
||||
let _user_id = current_user_id(&headers, &state)?;
|
||||
let value = crate::routes::admin::models_marketplace::marketplace_inner(&state.db).await?;
|
||||
let models = value
|
||||
.get("models")
|
||||
.and_then(Value::as_array)
|
||||
.map(|items| {
|
||||
items
|
||||
.iter()
|
||||
.map(|model| {
|
||||
let brand = model.get("name").and_then(Value::as_str).unwrap_or_default();
|
||||
json!({
|
||||
"name": brand,
|
||||
"accountCount": 0,
|
||||
"tokenCount": 0,
|
||||
"avgLatency": model.get("avgLatency").cloned().unwrap_or(Value::Null),
|
||||
"successRate": model.get("successRate").cloned().unwrap_or(Value::Null),
|
||||
"description": model.get("description").cloned().unwrap_or(Value::Null),
|
||||
"tags": model.get("tags").cloned().unwrap_or_else(|| json!([])),
|
||||
"supportedEndpointTypes": model.get("supportedEndpointTypes").cloned().unwrap_or_else(|| json!([])),
|
||||
"pricingSources": [],
|
||||
"accounts": [],
|
||||
})
|
||||
})
|
||||
.collect::<Vec<_>>()
|
||||
})
|
||||
.unwrap_or_default();
|
||||
|
||||
Ok(Json(json!({
|
||||
"models": models,
|
||||
"meta": value.get("meta").cloned().unwrap_or_else(|| json!({})),
|
||||
"userView": true,
|
||||
"includePricing": query.include_pricing.is_some(),
|
||||
})))
|
||||
}
|
||||
|
||||
@@ -1,785 +0,0 @@
|
||||
//! `/api/video-agent/*` — per-user multi-tenant refactor.
|
||||
//!
|
||||
//! Every handler resolves the caller from the JWT in the `Authorization`
|
||||
//! header (no anonymous access) and, where it makes sense, lazily starts
|
||||
//! that user's ComfyUI container via [`ComfyUIPool`]. The LLM/image/TTS
|
||||
//! helpers continue to call BoosAPI's own `/v1/*` loopback — they are
|
||||
//! independent of ComfyUI and stay shared.
|
||||
//!
|
||||
//! Status responses preserve every existing JSON key the frontend reads
|
||||
//! and add `containerState` so the new full-screen loading overlay can
|
||||
//! poll until the container is ready.
|
||||
|
||||
use axum::{
|
||||
body::Body,
|
||||
extract::State,
|
||||
http::{header, HeaderMap, Response, StatusCode},
|
||||
response::{
|
||||
sse::{Event, Sse},
|
||||
IntoResponse,
|
||||
},
|
||||
routing::{get, post},
|
||||
Json, Router,
|
||||
};
|
||||
use futures::stream;
|
||||
use serde_json::{json, Value};
|
||||
use std::convert::Infallible;
|
||||
use std::time::{Duration, Instant};
|
||||
use tracing::info;
|
||||
|
||||
use crate::{
|
||||
auth::jwt::validate_jwt,
|
||||
error::{AppError, AppResult},
|
||||
services::comfyui_pool::{ContainerState, PoolError},
|
||||
AppState,
|
||||
};
|
||||
|
||||
const COLD_START_HOLD: Duration = Duration::from_secs(30);
|
||||
const COLD_START_POLL: Duration = Duration::from_millis(500);
|
||||
const ADMIN_WORKSPACE_USER_ID: i64 = 0;
|
||||
|
||||
const SCRIPT_ANALYSIS_PROMPT: &str = r#"你是一个 AI 视频生成助手。分析用户提供的脚本，提取以下结构化信息，以 JSON 格式返回：
|
||||
|
||||
{
|
||||
"title": "视频标题",
|
||||
"characters": [
|
||||
{
|
||||
"name": "角色名",
|
||||
"gender": "男/女",
|
||||
"age": "年龄描述",
|
||||
"appearance": "外貌详细描述（发型、脸型、服装、体型等）",
|
||||
"personality": "性格特征"
|
||||
}
|
||||
],
|
||||
"scenes": [
|
||||
{
|
||||
"id": 1,
|
||||
"description": "场景描述",
|
||||
"environment": "环境详细描述（时间、地点、天气、光线等）",
|
||||
"characters": ["出现的角色名"],
|
||||
"action": "场景中发生的动作",
|
||||
"dialogue": [
|
||||
{"character": "角色名", "line": "台词", "emotion": "情感语气"}
|
||||
]
|
||||
}
|
||||
],
|
||||
"style": "整体视觉风格（动画/写实/像素等）"
|
||||
}
|
||||
|
||||
只返回 JSON，不要其他文字。"#;
|
||||
|
||||
const CHARACTER_PROMPT_GEN: &str = r#"根据角色描述生成 4 组稳定的图像生成提示词（正面、背面、左侧面、右侧面），用于 AI 图像生成模型。
|
||||
每组提示词应包含：人物姿态、服装、光照、背景。
|
||||
返回 JSON 格式：
|
||||
{
|
||||
"front": "正面提示词",
|
||||
"back": "背面提示词",
|
||||
"left": "左侧面提示词",
|
||||
"right": "右侧面提示词"
|
||||
}
|
||||
只返回 JSON。"#;
|
||||
|
||||
const SCENE_PROMPT_GEN: &str = r#"根据场景描述生成图像生成提示词，用于 AI 图像生成模型创建场景背景。
|
||||
包含：环境、光照、色调、氛围。
|
||||
返回 JSON：
|
||||
{
|
||||
"prompt": "场景生成提示词",
|
||||
"negative_prompt": "负面提示词"
|
||||
}
|
||||
只返回 JSON。"#;
|
||||
|
||||
pub fn routes() -> Router<AppState> {
|
||||
Router::new()
|
||||
.route("/api/video-agent/status", get(status))
|
||||
.route("/api/video-agent/start", post(start_comfyui))
|
||||
.route("/api/video-agent/stop", post(stop_comfyui))
|
||||
.route("/api/video-agent/restart", post(restart_comfyui))
|
||||
.route("/api/video-agent/chat", post(chat))
|
||||
.route("/api/video-agent/analyze", post(analyze))
|
||||
.route(
|
||||
"/api/video-agent/character-prompts",
|
||||
post(character_prompts),
|
||||
)
|
||||
.route(
|
||||
"/api/video-agent/generate-character",
|
||||
post(generate_character),
|
||||
)
|
||||
.route("/api/video-agent/scene-prompt", post(scene_prompt))
|
||||
.route("/api/video-agent/generate-scene", post(generate_scene))
|
||||
.route("/api/video-agent/generate-tts", post(generate_tts))
|
||||
.route("/api/video-agent/full-workflow", post(full_workflow))
|
||||
}
|
||||
|
||||
// ── Auth ───────────────────────────────────────────────────────────────
|
||||
|
||||
/// Resolve the caller workspace from `Authorization: Bearer <token>`.
|
||||
/// User JWTs map to their real user id; admin/global proxy tokens map to
|
||||
/// workspace 0 so the admin console can use the same Docker-backed ComfyUI.
|
||||
fn resolve_user_id(headers: &HeaderMap, state: &AppState) -> AppResult<i64> {
|
||||
let auth = headers
|
||||
.get(header::AUTHORIZATION)
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.and_then(|v| v.strip_prefix("Bearer "))
|
||||
.ok_or_else(|| AppError::Unauthorized("Missing or invalid Authorization header".into()))?;
|
||||
|
||||
if auth == state.config.auth.auth_token || auth == state.config.auth.proxy_token {
|
||||
return Ok(ADMIN_WORKSPACE_USER_ID);
|
||||
}
|
||||
|
||||
let claims = validate_jwt(&state.config.auth.jwt_secret, auth)
|
||||
.map_err(|_| AppError::Unauthorized("Invalid or expired token".into()))?;
|
||||
|
||||
claims
|
||||
.sub
|
||||
.parse::<i64>()
|
||||
.map_err(|_| AppError::Unauthorized("Invalid user ID in token".into()))
|
||||
}
|
||||
|
||||
// ── Pool helpers ───────────────────────────────────────────────────────
|
||||
|
||||
/// Map a [`PoolError`] to the appropriate [`AppError`].
|
||||
fn pool_err_to_app(e: PoolError) -> AppError {
|
||||
match e {
|
||||
PoolError::AtCapacity(_) => AppError::ServiceUnavailable("comfyui pool at capacity".into()),
|
||||
PoolError::StartTimeout(_) => {
|
||||
AppError::GatewayTimeout("comfyui container start timeout".into())
|
||||
}
|
||||
other => AppError::Internal(anyhow::anyhow!(other)),
|
||||
}
|
||||
}
|
||||
|
||||
/// Lazily start the caller's container and wait for readiness up to 30 s.
|
||||
/// Returns the host port for the loopback HTTP API.
|
||||
async fn ensure_container_ready(state: &AppState, user_id: i64) -> AppResult<u16> {
|
||||
state
|
||||
.comfyui_pool
|
||||
.get_or_start(user_id)
|
||||
.await
|
||||
.map_err(pool_err_to_app)?;
|
||||
|
||||
let deadline = Instant::now() + COLD_START_HOLD;
|
||||
loop {
|
||||
match state.comfyui_pool.state(user_id).await {
|
||||
ContainerState::Ready => break,
|
||||
ContainerState::Failed | ContainerState::Stopped => {
|
||||
return Err(AppError::UpstreamError(
|
||||
"comfyui container failed to start".into(),
|
||||
));
|
||||
}
|
||||
ContainerState::Starting => {
|
||||
if Instant::now() >= deadline {
|
||||
return Err(AppError::GatewayTimeout(
|
||||
"comfyui cold-start timeout".into(),
|
||||
));
|
||||
}
|
||||
tokio::time::sleep(COLD_START_POLL).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
state.comfyui_pool.touch(user_id).await;
|
||||
|
||||
state
|
||||
.comfyui_pool
|
||||
.entry(user_id)
|
||||
.await
|
||||
.map(|e| e.host_port)
|
||||
.ok_or_else(|| AppError::UpstreamError("container vanished after ready".into()))
|
||||
}
|
||||
|
||||
/// JSON status response. Preserves the keys the frontend already reads
|
||||
/// (`status`, `pid`, `uptime_secs`, `mode`, `model_backend`, `base_url`,
|
||||
/// `version`, `last_error`) and adds `containerState` (task 5.2).
|
||||
async fn pool_status_json(state: &AppState, user_id: i64) -> Value {
|
||||
let cs = state.comfyui_pool.state(user_id).await;
|
||||
let entry = state.comfyui_pool.entry(user_id).await;
|
||||
let port = entry.as_ref().map(|e| e.host_port);
|
||||
let last_error = entry.as_ref().and_then(|e| e.last_error.clone());
|
||||
|
||||
let status = match cs {
|
||||
ContainerState::Starting => "starting",
|
||||
ContainerState::Ready => "running",
|
||||
ContainerState::Stopped => "stopped",
|
||||
ContainerState::Failed => "unhealthy",
|
||||
};
|
||||
|
||||
let base_url = port
|
||||
.map(|p| format!("http://127.0.0.1:{}", p))
|
||||
.unwrap_or_default();
|
||||
|
||||
let model_backend = format!(
|
||||
"{}/v1",
|
||||
state
|
||||
.comfyui_pool
|
||||
.config()
|
||||
.boosapi_base
|
||||
.trim_end_matches("/v1")
|
||||
.trim_end_matches('/')
|
||||
);
|
||||
|
||||
json!({
|
||||
"status": status,
|
||||
"pid": Value::Null,
|
||||
"uptime_secs": 0,
|
||||
"mode": "container",
|
||||
"model_backend": model_backend,
|
||||
"base_url": base_url,
|
||||
"version": Value::Null,
|
||||
"last_error": last_error,
|
||||
"containerState": cs.as_str(),
|
||||
})
|
||||
}
|
||||
|
||||
// ── Lifecycle handlers ─────────────────────────────────────────────────
|
||||
|
||||
async fn status(State(state): State<AppState>, headers: HeaderMap) -> AppResult<Json<Value>> {
|
||||
let user_id = resolve_user_id(&headers, &state)?;
|
||||
Ok(Json(pool_status_json(&state, user_id).await))
|
||||
}
|
||||
|
||||
async fn start_comfyui(
|
||||
State(state): State<AppState>,
|
||||
headers: HeaderMap,
|
||||
) -> AppResult<Json<Value>> {
|
||||
let user_id = resolve_user_id(&headers, &state)?;
|
||||
// Idempotent: get_or_start returns the existing entry if already up.
|
||||
state
|
||||
.comfyui_pool
|
||||
.get_or_start(user_id)
|
||||
.await
|
||||
.map_err(pool_err_to_app)?;
|
||||
Ok(Json(json!({
|
||||
"success": true,
|
||||
"comfyui": pool_status_json(&state, user_id).await,
|
||||
})))
|
||||
}
|
||||
|
||||
async fn stop_comfyui(State(state): State<AppState>, headers: HeaderMap) -> AppResult<Json<Value>> {
|
||||
let user_id = resolve_user_id(&headers, &state)?;
|
||||
state
|
||||
.comfyui_pool
|
||||
.stop(user_id)
|
||||
.await
|
||||
.map_err(pool_err_to_app)?;
|
||||
Ok(Json(json!({
|
||||
"success": true,
|
||||
"comfyui": pool_status_json(&state, user_id).await,
|
||||
})))
|
||||
}
|
||||
|
||||
async fn restart_comfyui(
|
||||
State(state): State<AppState>,
|
||||
headers: HeaderMap,
|
||||
) -> AppResult<Json<Value>> {
|
||||
let user_id = resolve_user_id(&headers, &state)?;
|
||||
// Stop is best-effort: if the container is already gone, that's fine.
|
||||
let _ = state.comfyui_pool.stop(user_id).await;
|
||||
state
|
||||
.comfyui_pool
|
||||
.get_or_start(user_id)
|
||||
.await
|
||||
.map_err(pool_err_to_app)?;
|
||||
Ok(Json(json!({
|
||||
"success": true,
|
||||
"comfyui": pool_status_json(&state, user_id).await,
|
||||
})))
|
||||
}
|
||||
|
||||
// ── BoosAPI loopback helpers ───────────────────────────────────────────
|
||||
//
|
||||
// LLM, image, and TTS calls go to BoosAPI's own `/v1/*` endpoints — they
|
||||
// are not ComfyUI traffic. We pull the base URL from the same config the
|
||||
// pool uses so there's exactly one source of truth.
|
||||
|
||||
fn boosapi_v1_base(state: &AppState) -> String {
|
||||
let raw = state
|
||||
.comfyui_pool
|
||||
.config()
|
||||
.boosapi_base
|
||||
.trim_end_matches('/');
|
||||
if raw.ends_with("/v1") {
|
||||
raw.to_string()
|
||||
} else {
|
||||
format!("{}/v1", raw)
|
||||
}
|
||||
}
|
||||
|
||||
async fn chat(
|
||||
State(state): State<AppState>,
|
||||
headers: HeaderMap,
|
||||
Json(body): Json<Value>,
|
||||
) -> AppResult<Sse<impl futures::Stream<Item = Result<Event, Infallible>>>> {
|
||||
let user_id = resolve_user_id(&headers, &state)?;
|
||||
// Lazily warm the container so a subsequent generate-* call doesn't
|
||||
// pay the full cold start.
|
||||
let _ = state.comfyui_pool.get_or_start(user_id).await;
|
||||
state.comfyui_pool.touch(user_id).await;
|
||||
|
||||
let prompt = body
|
||||
.get("message")
|
||||
.and_then(Value::as_str)
|
||||
.or_else(|| {
|
||||
body.get("messages")
|
||||
.and_then(Value::as_array)
|
||||
.and_then(|messages| messages.last())
|
||||
.and_then(|message| message.get("content"))
|
||||
.and_then(Value::as_str)
|
||||
})
|
||||
.unwrap_or("生成一个短视频工作流");
|
||||
|
||||
let workflow = build_basic_workflow(prompt);
|
||||
let events = vec![
|
||||
json!({
|
||||
"type": "text",
|
||||
"content": "已收到需求，正在生成 ComfyUI workflow..."
|
||||
}),
|
||||
json!({
|
||||
"type": "workflow_generated",
|
||||
"workflow": workflow
|
||||
}),
|
||||
json!({
|
||||
"type": "done"
|
||||
}),
|
||||
];
|
||||
|
||||
Ok(Sse::new(stream::iter(events.into_iter().map(|event| {
|
||||
Ok(Event::default().data(event.to_string()))
|
||||
}))))
|
||||
}
|
||||
|
||||
fn build_basic_workflow(prompt: &str) -> Value {
|
||||
json!({
|
||||
"boosapi": {
|
||||
"kind": "video-agent-workflow",
|
||||
"version": 1,
|
||||
"prompt": prompt,
|
||||
},
|
||||
"nodes": {},
|
||||
})
|
||||
}
|
||||
|
||||
fn required_str<'a>(body: &'a Value, key: &str) -> AppResult<&'a str> {
|
||||
body.get(key)
|
||||
.and_then(Value::as_str)
|
||||
.filter(|s| !s.trim().is_empty())
|
||||
.ok_or_else(|| AppError::BadRequest(format!("{key} is required")))
|
||||
}
|
||||
|
||||
fn required_value<'a>(body: &'a Value, key: &str) -> AppResult<&'a Value> {
|
||||
body.get(key)
|
||||
.filter(|v| !v.is_null())
|
||||
.ok_or_else(|| AppError::BadRequest(format!("{key} is required")))
|
||||
}
|
||||
|
||||
fn strip_json_markers(s: &str) -> String {
|
||||
let mut text = s.trim().to_string();
|
||||
|
||||
if text.starts_with("```json") {
|
||||
text = text.trim_start_matches("```json").to_string();
|
||||
} else if text.starts_with("```") {
|
||||
text = text.trim_start_matches("```").to_string();
|
||||
}
|
||||
|
||||
if text.ends_with("```") {
|
||||
text = text.trim_end_matches("```").to_string();
|
||||
}
|
||||
|
||||
text.trim().to_string()
|
||||
}
|
||||
|
||||
fn parse_llm_json(raw: &str) -> AppResult<Value> {
|
||||
let cleaned = strip_json_markers(raw);
|
||||
serde_json::from_str::<Value>(&cleaned).map_err(|e| {
|
||||
AppError::Internal(anyhow::anyhow!(
|
||||
"failed to parse LLM JSON response: {}; response: {}",
|
||||
e,
|
||||
cleaned
|
||||
))
|
||||
})
|
||||
}
|
||||
|
||||
async fn call_llm(state: &AppState, prompt: &str, system: &str) -> AppResult<String> {
|
||||
let client = reqwest::Client::new();
|
||||
let url = format!("{}/chat/completions", boosapi_v1_base(state));
|
||||
|
||||
info!("video-agent calling BoosAPI LLM endpoint: {}", url);
|
||||
|
||||
let resp = client
|
||||
.post(url)
|
||||
.bearer_auth(&state.config.auth.proxy_token)
|
||||
.json(&json!({
|
||||
"model": &state.config.comfyui.agent_model,
|
||||
"messages": [
|
||||
{ "role": "system", "content": system },
|
||||
{ "role": "user", "content": prompt }
|
||||
],
|
||||
"temperature": 0.7
|
||||
}))
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| AppError::Internal(anyhow::anyhow!("LLM request failed: {}", e)))?;
|
||||
|
||||
let status = resp.status();
|
||||
let data: Value = resp
|
||||
.json()
|
||||
.await
|
||||
.map_err(|e| AppError::Internal(anyhow::anyhow!("LLM response parse failed: {}", e)))?;
|
||||
|
||||
if !status.is_success() {
|
||||
return Err(AppError::Internal(anyhow::anyhow!(
|
||||
"LLM upstream returned {}: {}",
|
||||
status,
|
||||
data
|
||||
)));
|
||||
}
|
||||
|
||||
Ok(data
|
||||
.get("choices")
|
||||
.and_then(Value::as_array)
|
||||
.and_then(|choices| choices.first())
|
||||
.and_then(|choice| choice.get("message"))
|
||||
.and_then(|message| message.get("content"))
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or("")
|
||||
.to_string())
|
||||
}
|
||||
|
||||
async fn call_image_gen(state: &AppState, prompt: &str) -> AppResult<String> {
|
||||
let client = reqwest::Client::new();
|
||||
let url = format!("{}/images/generations", boosapi_v1_base(state));
|
||||
|
||||
info!("video-agent calling BoosAPI image endpoint: {}", url);
|
||||
|
||||
let resp = client
|
||||
.post(url)
|
||||
.bearer_auth(&state.config.auth.proxy_token)
|
||||
.json(&json!({
|
||||
"model": "gpt-image-2",
|
||||
"prompt": prompt,
|
||||
"n": 1,
|
||||
"size": "1024x1024"
|
||||
}))
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| AppError::Internal(anyhow::anyhow!("image request failed: {}", e)))?;
|
||||
|
||||
let status = resp.status();
|
||||
let data: Value = resp
|
||||
.json()
|
||||
.await
|
||||
.map_err(|e| AppError::Internal(anyhow::anyhow!("image response parse failed: {}", e)))?;
|
||||
|
||||
if !status.is_success() {
|
||||
return Err(AppError::Internal(anyhow::anyhow!(
|
||||
"image upstream returned {}: {}",
|
||||
status,
|
||||
data
|
||||
)));
|
||||
}
|
||||
|
||||
Ok(data
|
||||
.get("data")
|
||||
.and_then(Value::as_array)
|
||||
.and_then(|items| items.first())
|
||||
.and_then(|item| item.get("url"))
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or("")
|
||||
.to_string())
|
||||
}
|
||||
|
||||
async fn call_tts(state: &AppState, text: &str, voice: &str) -> AppResult<Vec<u8>> {
|
||||
let client = reqwest::Client::new();
|
||||
let url = format!("{}/audio/speech", boosapi_v1_base(state));
|
||||
|
||||
info!("video-agent calling BoosAPI TTS endpoint: {}", url);
|
||||
|
||||
let resp = client
|
||||
.post(url)
|
||||
.bearer_auth(&state.config.auth.proxy_token)
|
||||
.json(&json!({
|
||||
"model": "tts-1",
|
||||
"input": text,
|
||||
"voice": voice
|
||||
}))
|
||||
.send()
|
||||
.await
|
||||
.map_err(|e| AppError::Internal(anyhow::anyhow!("TTS request failed: {}", e)))?;
|
||||
|
||||
let status = resp.status();
|
||||
|
||||
if !status.is_success() {
|
||||
let body = resp.text().await.unwrap_or_default();
|
||||
return Err(AppError::Internal(anyhow::anyhow!(
|
||||
"TTS upstream returned {}: {}",
|
||||
status,
|
||||
body
|
||||
)));
|
||||
}
|
||||
|
||||
let bytes = resp
|
||||
.bytes()
|
||||
.await
|
||||
.map_err(|e| AppError::Internal(anyhow::anyhow!("TTS response read failed: {}", e)))?;
|
||||
|
||||
Ok(bytes.to_vec())
|
||||
}
|
||||
|
||||
// ── Per-request handlers (5.5) ─────────────────────────────────────────
|
||||
|
||||
async fn analyze(
|
||||
State(state): State<AppState>,
|
||||
headers: HeaderMap,
|
||||
Json(body): Json<Value>,
|
||||
) -> AppResult<Json<Value>> {
|
||||
let user_id = resolve_user_id(&headers, &state)?;
|
||||
state.comfyui_pool.touch(user_id).await;
|
||||
|
||||
let script = required_str(&body, "script")?;
|
||||
|
||||
let result = call_llm(&state, script, SCRIPT_ANALYSIS_PROMPT).await?;
|
||||
let parsed = parse_llm_json(&result)?;
|
||||
|
||||
Ok(Json(json!({
|
||||
"success": true,
|
||||
"data": parsed
|
||||
})))
|
||||
}
|
||||
|
||||
async fn character_prompts(
|
||||
State(state): State<AppState>,
|
||||
headers: HeaderMap,
|
||||
Json(body): Json<Value>,
|
||||
) -> AppResult<Json<Value>> {
|
||||
let user_id = resolve_user_id(&headers, &state)?;
|
||||
state.comfyui_pool.touch(user_id).await;
|
||||
|
||||
let character = required_value(&body, "character")?;
|
||||
|
||||
let name = character
|
||||
.get("name")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default();
|
||||
let appearance = character
|
||||
.get("appearance")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default();
|
||||
let gender = character
|
||||
.get("gender")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default();
|
||||
let age = character
|
||||
.get("age")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default();
|
||||
|
||||
let desc = format!("{name}: {appearance}, {gender}, {age}");
|
||||
let result = call_llm(&state, &desc, CHARACTER_PROMPT_GEN).await?;
|
||||
let prompts = parse_llm_json(&result)?;
|
||||
|
||||
Ok(Json(json!({
|
||||
"success": true,
|
||||
"data": prompts
|
||||
})))
|
||||
}
|
||||
|
||||
async fn generate_character(
|
||||
State(state): State<AppState>,
|
||||
headers: HeaderMap,
|
||||
Json(body): Json<Value>,
|
||||
) -> AppResult<Json<Value>> {
|
||||
let user_id = resolve_user_id(&headers, &state)?;
|
||||
// Image generation can take a while; ensure the container is warm so
|
||||
// downstream ComfyUI workflow reuse doesn't pay another cold start.
|
||||
let _ = ensure_container_ready(&state, user_id).await?;
|
||||
|
||||
let prompts = required_value(&body, "prompts")?;
|
||||
let prompt_obj = prompts
|
||||
.as_object()
|
||||
.ok_or_else(|| AppError::BadRequest("prompts must be an object".to_string()))?;
|
||||
|
||||
let mut results = serde_json::Map::new();
|
||||
|
||||
for (angle, prompt_value) in prompt_obj {
|
||||
let prompt = prompt_value.as_str().ok_or_else(|| {
|
||||
AppError::BadRequest(format!("prompt for angle '{angle}' must be a string"))
|
||||
})?;
|
||||
|
||||
let url = call_image_gen(&state, prompt).await?;
|
||||
results.insert(angle.clone(), Value::String(url));
|
||||
}
|
||||
|
||||
state.comfyui_pool.touch(user_id).await;
|
||||
|
||||
Ok(Json(json!({
|
||||
"success": true,
|
||||
"data": Value::Object(results)
|
||||
})))
|
||||
}
|
||||
|
||||
async fn scene_prompt(
|
||||
State(state): State<AppState>,
|
||||
headers: HeaderMap,
|
||||
Json(body): Json<Value>,
|
||||
) -> AppResult<Json<Value>> {
|
||||
let user_id = resolve_user_id(&headers, &state)?;
|
||||
state.comfyui_pool.touch(user_id).await;
|
||||
|
||||
let scene = required_value(&body, "scene")?;
|
||||
|
||||
let environment = scene
|
||||
.get("environment")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default();
|
||||
let description = scene
|
||||
.get("description")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default();
|
||||
|
||||
let desc = format!("{environment}. {description}");
|
||||
let result = call_llm(&state, &desc, SCENE_PROMPT_GEN).await?;
|
||||
let prompt = parse_llm_json(&result)?;
|
||||
|
||||
Ok(Json(json!({
|
||||
"success": true,
|
||||
"data": prompt
|
||||
})))
|
||||
}
|
||||
|
||||
async fn generate_scene(
|
||||
State(state): State<AppState>,
|
||||
headers: HeaderMap,
|
||||
Json(body): Json<Value>,
|
||||
) -> AppResult<Json<Value>> {
|
||||
let user_id = resolve_user_id(&headers, &state)?;
|
||||
let _ = ensure_container_ready(&state, user_id).await?;
|
||||
|
||||
let prompt = required_str(&body, "prompt")?;
|
||||
let url = call_image_gen(&state, prompt).await?;
|
||||
|
||||
state.comfyui_pool.touch(user_id).await;
|
||||
|
||||
Ok(Json(json!({
|
||||
"success": true,
|
||||
"data": {
|
||||
"url": url
|
||||
}
|
||||
})))
|
||||
}
|
||||
|
||||
async fn generate_tts(
|
||||
State(state): State<AppState>,
|
||||
headers: HeaderMap,
|
||||
Json(body): Json<Value>,
|
||||
) -> AppResult<Response<Body>> {
|
||||
let user_id = resolve_user_id(&headers, &state)?;
|
||||
state.comfyui_pool.touch(user_id).await;
|
||||
|
||||
let text = required_str(&body, "text")?;
|
||||
let voice = body
|
||||
.get("voice")
|
||||
.and_then(Value::as_str)
|
||||
.filter(|s| !s.trim().is_empty())
|
||||
.unwrap_or("alloy");
|
||||
|
||||
let audio = call_tts(&state, text, voice).await?;
|
||||
|
||||
Ok((
|
||||
StatusCode::OK,
|
||||
[(header::CONTENT_TYPE, "audio/mpeg")],
|
||||
audio,
|
||||
)
|
||||
.into_response())
|
||||
}
|
||||
|
||||
async fn full_workflow(
|
||||
State(state): State<AppState>,
|
||||
headers: HeaderMap,
|
||||
Json(body): Json<Value>,
|
||||
) -> AppResult<Json<Value>> {
|
||||
let user_id = resolve_user_id(&headers, &state)?;
|
||||
let _ = ensure_container_ready(&state, user_id).await?;
|
||||
|
||||
let script = required_str(&body, "script")?;
|
||||
let requested_style = body.get("style").and_then(Value::as_str);
|
||||
|
||||
let analysis_raw = call_llm(&state, script, SCRIPT_ANALYSIS_PROMPT).await?;
|
||||
let analysis = parse_llm_json(&analysis_raw)?;
|
||||
|
||||
let mut characters = Vec::new();
|
||||
|
||||
if let Some(chars) = analysis.get("characters").and_then(Value::as_array) {
|
||||
for character in chars {
|
||||
let name = character
|
||||
.get("name")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default();
|
||||
let appearance = character
|
||||
.get("appearance")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default();
|
||||
let gender = character
|
||||
.get("gender")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default();
|
||||
let age = character
|
||||
.get("age")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default();
|
||||
|
||||
let desc = format!("{name}: {appearance}, {gender}, {age}");
|
||||
let prompts_raw = call_llm(&state, &desc, CHARACTER_PROMPT_GEN).await?;
|
||||
let prompts = parse_llm_json(&prompts_raw)?;
|
||||
|
||||
let mut item = character.clone();
|
||||
if let Some(obj) = item.as_object_mut() {
|
||||
obj.insert("imagePrompts".to_string(), prompts);
|
||||
}
|
||||
|
||||
characters.push(item);
|
||||
}
|
||||
}
|
||||
|
||||
let mut scenes = Vec::new();
|
||||
|
||||
if let Some(scene_items) = analysis.get("scenes").and_then(Value::as_array) {
|
||||
for scene in scene_items {
|
||||
let environment = scene
|
||||
.get("environment")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default();
|
||||
let description = scene
|
||||
.get("description")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default();
|
||||
|
||||
let desc = format!("{environment}. {description}");
|
||||
let scene_prompt_raw = call_llm(&state, &desc, SCENE_PROMPT_GEN).await?;
|
||||
let image_prompt = parse_llm_json(&scene_prompt_raw)?;
|
||||
|
||||
let mut item = scene.clone();
|
||||
if let Some(obj) = item.as_object_mut() {
|
||||
obj.insert("imagePrompt".to_string(), image_prompt);
|
||||
}
|
||||
|
||||
scenes.push(item);
|
||||
}
|
||||
}
|
||||
|
||||
let style = analysis
|
||||
.get("style")
|
||||
.cloned()
|
||||
.or_else(|| requested_style.map(|s| Value::String(s.to_string())))
|
||||
.unwrap_or(Value::Null);
|
||||
|
||||
state.comfyui_pool.touch(user_id).await;
|
||||
|
||||
Ok(Json(json!({
|
||||
"success": true,
|
||||
"data": {
|
||||
"title": analysis.get("title").cloned().unwrap_or(Value::Null),
|
||||
"style": style,
|
||||
"characters": characters,
|
||||
"scenes": scenes
|
||||
}
|
||||
})))
|
||||
}
|
||||
@@ -93,10 +93,7 @@ pub fn routes() -> Router<AppState> {
|
||||
|
||||
async fn status(State(state): State<AppState>, headers: HeaderMap) -> AppResult<Json<Value>> {
|
||||
check_auth(&headers, &state)?;
|
||||
Ok(Json(
|
||||
serde_json::to_value(state.comfyui.status_response().await)
|
||||
.unwrap_or_else(|_| json!({ "status": "unknown" })),
|
||||
))
|
||||
Ok(Json(external_comfyui_status(&state).await))
|
||||
}
|
||||
|
||||
async fn start_comfyui(
|
||||
@@ -104,27 +101,19 @@ async fn start_comfyui(
|
||||
headers: HeaderMap,
|
||||
) -> AppResult<Json<Value>> {
|
||||
check_auth(&headers, &state)?;
|
||||
state
|
||||
.comfyui
|
||||
.start()
|
||||
.await
|
||||
.map_err(|err| AppError::Internal(anyhow::anyhow!(err)))?;
|
||||
Ok(Json(json!({
|
||||
"success": true,
|
||||
"comfyui": state.comfyui.status_response().await
|
||||
"message": "内置 Docker ComfyUI 工作台已移除；请配置外部 ComfyUI 或使用 /workflows 工作流画布。",
|
||||
"comfyui": external_comfyui_status_value(&state).await
|
||||
})))
|
||||
}
|
||||
|
||||
async fn stop_comfyui(State(state): State<AppState>, headers: HeaderMap) -> AppResult<Json<Value>> {
|
||||
check_auth(&headers, &state)?;
|
||||
state
|
||||
.comfyui
|
||||
.stop()
|
||||
.await
|
||||
.map_err(|err| AppError::Internal(anyhow::anyhow!(err)))?;
|
||||
Ok(Json(json!({
|
||||
"success": true,
|
||||
"comfyui": state.comfyui.status_response().await
|
||||
"message": "内置 Docker ComfyUI 工作台已移除；BoosAPI 不再管理 ComfyUI 进程生命周期。",
|
||||
"comfyui": external_comfyui_status_value(&state).await
|
||||
})))
|
||||
}
|
||||
|
||||
@@ -133,14 +122,10 @@ async fn restart_comfyui(
|
||||
headers: HeaderMap,
|
||||
) -> AppResult<Json<Value>> {
|
||||
check_auth(&headers, &state)?;
|
||||
state
|
||||
.comfyui
|
||||
.restart()
|
||||
.await
|
||||
.map_err(|err| AppError::Internal(anyhow::anyhow!(err)))?;
|
||||
Ok(Json(json!({
|
||||
"success": true,
|
||||
"comfyui": state.comfyui.status_response().await
|
||||
"message": "内置 Docker ComfyUI 工作台已移除；请在外部 ComfyUI 服务侧重启。",
|
||||
"comfyui": external_comfyui_status_value(&state).await
|
||||
})))
|
||||
}
|
||||
|
||||
@@ -187,6 +172,33 @@ fn local_api_base(state: &AppState) -> String {
|
||||
state.comfyui.model_backend_url()
|
||||
}
|
||||
|
||||
async fn external_comfyui_status(state: &AppState) -> Value {
|
||||
json!({
|
||||
"success": true,
|
||||
"comfyui": external_comfyui_status_value(state).await,
|
||||
})
|
||||
}
|
||||
|
||||
async fn external_comfyui_status_value(state: &AppState) -> Value {
|
||||
let status = state.comfyui.status_response().await;
|
||||
let reachable = state.comfyui.health_check().await;
|
||||
json!({
|
||||
"status": if reachable { "running" } else { "unavailable" },
|
||||
"pid": null,
|
||||
"uptime_secs": 0,
|
||||
"mode": "external",
|
||||
"model_backend": status.model_backend,
|
||||
"base_url": status.base_url,
|
||||
"version": status.version,
|
||||
"last_error": if reachable {
|
||||
Value::Null
|
||||
} else {
|
||||
json!("外部 ComfyUI 不可达；内置 Docker ComfyUI 工作台已移除，请配置 comfyui.base_url 或使用 /workflows 工作流画布。")
|
||||
},
|
||||
"dockerManaged": false,
|
||||
})
|
||||
}
|
||||
|
||||
fn check_auth(headers: &HeaderMap, state: &AppState) -> AppResult<()> {
|
||||
let expected = format!("Bearer {}", state.config.auth.auth_token);
|
||||
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
+31
-6
@@ -36,11 +36,23 @@ struct AssetQuery {
|
||||
pub fn routes() -> Router<AppState> {
|
||||
Router::new()
|
||||
.route("/api/workspace/assets", get(list_admin_assets))
|
||||
.route("/api/workspace/assets/{public_id}", get(get_admin_asset).delete(delete_admin_asset))
|
||||
.route("/api/workspace/assets/{public_id}/file", get(get_admin_asset_file))
|
||||
.route(
|
||||
"/api/workspace/assets/{public_id}",
|
||||
get(get_admin_asset).delete(delete_admin_asset),
|
||||
)
|
||||
.route(
|
||||
"/api/workspace/assets/{public_id}/file",
|
||||
get(get_admin_asset_file),
|
||||
)
|
||||
.route("/api/user/workspace/assets", get(list_user_assets))
|
||||
.route("/api/user/workspace/assets/{public_id}", get(get_user_asset).delete(delete_user_asset))
|
||||
.route("/api/user/workspace/assets/{public_id}/file", get(get_user_asset_file))
|
||||
.route(
|
||||
"/api/user/workspace/assets/{public_id}",
|
||||
get(get_user_asset).delete(delete_user_asset),
|
||||
)
|
||||
.route(
|
||||
"/api/user/workspace/assets/{public_id}/file",
|
||||
get(get_user_asset_file),
|
||||
)
|
||||
}
|
||||
|
||||
fn bearer(headers: &HeaderMap) -> Option<&str> {
|
||||
@@ -56,7 +68,11 @@ fn bearer(headers: &HeaderMap) -> Option<&str> {
|
||||
.filter(|value| !value.is_empty())
|
||||
}
|
||||
|
||||
fn resolve_caller(headers: &HeaderMap, state: &AppState, admin_required: bool) -> AppResult<WorkspaceCaller> {
|
||||
fn resolve_caller(
|
||||
headers: &HeaderMap,
|
||||
state: &AppState,
|
||||
admin_required: bool,
|
||||
) -> AppResult<WorkspaceCaller> {
|
||||
let token = bearer(headers)
|
||||
.ok_or_else(|| AppError::Unauthorized("Missing Authorization header".into()))?;
|
||||
|
||||
@@ -79,7 +95,16 @@ fn resolve_caller(headers: &HeaderMap, state: &AppState, admin_required: bool) -
|
||||
Err(AppError::Forbidden("Workspace access denied".into()))
|
||||
}
|
||||
|
||||
fn normalize_query(query: &AssetQuery) -> (i64, i64, Option<String>, Option<String>, Option<String>, Option<i64>) {
|
||||
fn normalize_query(
|
||||
query: &AssetQuery,
|
||||
) -> (
|
||||
i64,
|
||||
i64,
|
||||
Option<String>,
|
||||
Option<String>,
|
||||
Option<String>,
|
||||
Option<i64>,
|
||||
) {
|
||||
let page_size = query.page_size.unwrap_or(36).clamp(1, 120);
|
||||
let page = query.page.unwrap_or(1).max(1);
|
||||
let offset = (page - 1) * page_size;
|
||||
|
||||
@@ -1,672 +0,0 @@
|
||||
//! Per-user ComfyUI container pool.
|
||||
//!
|
||||
//! Owns the in-memory map `user_id -> PoolEntry`, allocates ports from the
|
||||
//! 8200..=8299 range, enforces the 30-container ceiling, runs the readiness
|
||||
//! probe, and reaps idle containers.
|
||||
//!
|
||||
//! All Docker/K3S specifics are encapsulated by the [`ContainerRuntime`]
|
||||
//! trait; this file MUST stay backend-agnostic.
|
||||
|
||||
use std::collections::{HashMap, HashSet};
|
||||
use std::path::PathBuf;
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use parking_lot::Mutex;
|
||||
use tokio::sync::{Mutex as AsyncMutex, RwLock};
|
||||
use tracing::{debug, error, info, warn};
|
||||
|
||||
use super::container_runtime::{
|
||||
BindMount, ContainerRuntime, ContainerSpec, ResourceLimits, RuntimeError, TmpfsMount,
|
||||
};
|
||||
|
||||
/// Maximum concurrent containers (§design Decision 2: 30 × 2 GB ≈ 60 GB).
|
||||
pub const DEFAULT_MAX_CONCURRENT: usize = 30;
|
||||
/// Idle timeout before the reaper stops a container.
|
||||
pub const DEFAULT_IDLE_TIMEOUT: Duration = Duration::from_secs(60 * 60);
|
||||
/// Host port range for per-user containers.
|
||||
pub const PORT_RANGE_START: u16 = 8200;
|
||||
pub const PORT_RANGE_END: u16 = 8299;
|
||||
/// Time we wait for a freshly started container to answer `/system_stats`.
|
||||
pub const READINESS_TIMEOUT: Duration = Duration::from_secs(30);
|
||||
/// How often the reaper task wakes up.
|
||||
pub const REAPER_TICK: Duration = Duration::from_secs(60);
|
||||
/// Image tag we run for per-user tenants.
|
||||
pub const DEFAULT_IMAGE: &str = "boosapi-comfyui:multi-tenant";
|
||||
/// Host directory under which per-user volumes live.
|
||||
pub const DEFAULT_DATA_ROOT: &str = "/data/comfyui-users";
|
||||
|
||||
/// Lifecycle state visible to callers (frontend renders loading vs ready).
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum ContainerState {
|
||||
Starting,
|
||||
Ready,
|
||||
Stopped,
|
||||
Failed,
|
||||
}
|
||||
|
||||
impl ContainerState {
|
||||
pub fn as_str(self) -> &'static str {
|
||||
match self {
|
||||
ContainerState::Starting => "starting",
|
||||
ContainerState::Ready => "ready",
|
||||
ContainerState::Stopped => "stopped",
|
||||
ContainerState::Failed => "failed",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// One running (or starting) per-user container.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PoolEntry {
|
||||
pub user_id: i64,
|
||||
pub container_name: String,
|
||||
pub host_port: u16,
|
||||
pub state: ContainerState,
|
||||
pub last_used_at: Instant,
|
||||
pub active_ws: u32,
|
||||
pub last_error: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(thiserror::Error, Debug)]
|
||||
pub enum PoolError {
|
||||
#[error("pool at capacity ({0} containers running)")]
|
||||
AtCapacity(usize),
|
||||
#[error("no free port in pool range")]
|
||||
OutOfPorts,
|
||||
#[error("container failed to become ready in {0:?}")]
|
||||
StartTimeout(Duration),
|
||||
#[error("runtime error: {0}")]
|
||||
Runtime(#[from] RuntimeError),
|
||||
#[error("internal: {0}")]
|
||||
Internal(String),
|
||||
}
|
||||
|
||||
pub type PoolResult<T> = Result<T, PoolError>;
|
||||
|
||||
/// Tunable knobs the binary wires up at boot.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PoolConfig {
|
||||
pub image: String,
|
||||
pub max_concurrent: usize,
|
||||
pub idle_timeout: Duration,
|
||||
pub data_root: PathBuf,
|
||||
pub boosapi_base: String,
|
||||
pub builtin_key: String,
|
||||
}
|
||||
|
||||
impl PoolConfig {
|
||||
pub fn from_env(boosapi_base: String, builtin_key: String) -> Self {
|
||||
let image = std::env::var("BOOSAPI_COMFY_IMAGE").unwrap_or_else(|_| DEFAULT_IMAGE.into());
|
||||
let max_concurrent: usize = std::env::var("BOOSAPI_COMFY_MAX")
|
||||
.ok()
|
||||
.and_then(|v| v.parse().ok())
|
||||
.unwrap_or(DEFAULT_MAX_CONCURRENT);
|
||||
let idle_secs: u64 = std::env::var("BOOSAPI_COMFY_IDLE_SECS")
|
||||
.ok()
|
||||
.and_then(|v| v.parse().ok())
|
||||
.unwrap_or(DEFAULT_IDLE_TIMEOUT.as_secs());
|
||||
let data_root = std::env::var("BOOSAPI_COMFY_DATA_ROOT")
|
||||
.map(PathBuf::from)
|
||||
.unwrap_or_else(|_| PathBuf::from(DEFAULT_DATA_ROOT));
|
||||
Self {
|
||||
image,
|
||||
max_concurrent,
|
||||
idle_timeout: Duration::from_secs(idle_secs),
|
||||
data_root,
|
||||
boosapi_base,
|
||||
builtin_key,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Thread-safe allocator over `PORT_RANGE_START..=PORT_RANGE_END`.
|
||||
struct PortAllocator {
|
||||
in_use: Mutex<HashSet<u16>>,
|
||||
}
|
||||
|
||||
impl PortAllocator {
|
||||
fn new() -> Self {
|
||||
Self {
|
||||
in_use: Mutex::new(HashSet::new()),
|
||||
}
|
||||
}
|
||||
|
||||
fn claim(&self, port: u16) {
|
||||
self.in_use.lock().insert(port);
|
||||
}
|
||||
|
||||
fn allocate(&self) -> Option<u16> {
|
||||
let mut g = self.in_use.lock();
|
||||
for p in PORT_RANGE_START..=PORT_RANGE_END {
|
||||
if !g.contains(&p) {
|
||||
g.insert(p);
|
||||
return Some(p);
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
fn release(&self, port: u16) {
|
||||
self.in_use.lock().remove(&port);
|
||||
}
|
||||
}
|
||||
|
||||
/// Per-user start lock. Prevents two concurrent calls to
|
||||
/// `get_or_start(7)` from spawning two containers for user 7.
|
||||
type StartLock = Arc<AsyncMutex<()>>;
|
||||
|
||||
/// The pool itself.
|
||||
pub struct ComfyUIPool {
|
||||
runtime: Arc<dyn ContainerRuntime>,
|
||||
config: PoolConfig,
|
||||
containers: RwLock<HashMap<i64, PoolEntry>>,
|
||||
start_locks: Mutex<HashMap<i64, StartLock>>,
|
||||
ports: PortAllocator,
|
||||
http: reqwest::Client,
|
||||
}
|
||||
|
||||
impl ComfyUIPool {
|
||||
pub fn new(runtime: Arc<dyn ContainerRuntime>, config: PoolConfig) -> Arc<Self> {
|
||||
Arc::new(Self {
|
||||
runtime,
|
||||
config,
|
||||
containers: RwLock::new(HashMap::new()),
|
||||
start_locks: Mutex::new(HashMap::new()),
|
||||
ports: PortAllocator::new(),
|
||||
http: reqwest::Client::builder()
|
||||
.timeout(Duration::from_secs(5))
|
||||
.build()
|
||||
.expect("reqwest client"),
|
||||
})
|
||||
}
|
||||
|
||||
pub fn config(&self) -> &PoolConfig {
|
||||
&self.config
|
||||
}
|
||||
|
||||
/// Start the user's container if it isn't already running and return a
|
||||
/// handle. Idempotent.
|
||||
pub async fn get_or_start(&self, user_id: i64) -> PoolResult<PoolEntry> {
|
||||
// Fast path: already ready.
|
||||
let existing = {
|
||||
let map = self.containers.read().await;
|
||||
map.get(&user_id).cloned()
|
||||
};
|
||||
if let Some(entry) = existing {
|
||||
if entry.state == ContainerState::Ready {
|
||||
self.touch_inner(user_id).await;
|
||||
return Ok(entry);
|
||||
}
|
||||
}
|
||||
|
||||
// Per-user start lock — only one task gets to spawn per user.
|
||||
let lock = self
|
||||
.start_locks
|
||||
.lock()
|
||||
.entry(user_id)
|
||||
.or_insert_with(|| Arc::new(AsyncMutex::new(())))
|
||||
.clone();
|
||||
let _g = lock.lock().await;
|
||||
|
||||
// Re-check after acquiring the lock; another task may have done the work.
|
||||
let existing = {
|
||||
let map = self.containers.read().await;
|
||||
map.get(&user_id).cloned()
|
||||
};
|
||||
if let Some(entry) = existing {
|
||||
if entry.state == ContainerState::Ready {
|
||||
self.touch_inner(user_id).await;
|
||||
return Ok(entry);
|
||||
}
|
||||
if entry.state == ContainerState::Starting {
|
||||
// Probe ourselves — the previous winner may have died mid-start.
|
||||
return self.probe_until_ready(user_id, entry).await;
|
||||
}
|
||||
}
|
||||
|
||||
// Capacity check.
|
||||
{
|
||||
let map = self.containers.read().await;
|
||||
let live = map
|
||||
.values()
|
||||
.filter(|e| matches!(e.state, ContainerState::Starting | ContainerState::Ready))
|
||||
.count();
|
||||
if live >= self.config.max_concurrent {
|
||||
return Err(PoolError::AtCapacity(live));
|
||||
}
|
||||
}
|
||||
|
||||
// Allocate port.
|
||||
let port = self.ports.allocate().ok_or(PoolError::OutOfPorts)?;
|
||||
let name = format!("comfyui-user-{}", user_id);
|
||||
|
||||
// Insert a Starting marker so concurrent callers see it.
|
||||
{
|
||||
let mut map = self.containers.write().await;
|
||||
map.insert(
|
||||
user_id,
|
||||
PoolEntry {
|
||||
user_id,
|
||||
container_name: name.clone(),
|
||||
host_port: port,
|
||||
state: ContainerState::Starting,
|
||||
last_used_at: Instant::now(),
|
||||
active_ws: 0,
|
||||
last_error: None,
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
let spec = self.build_spec(user_id, &name, port);
|
||||
|
||||
match self.runtime.start(spec).await {
|
||||
Ok(handle) => {
|
||||
debug!(user_id, container = %handle.name, port = handle.host_port, "container started");
|
||||
}
|
||||
Err(e) => {
|
||||
self.ports.release(port);
|
||||
let mut map = self.containers.write().await;
|
||||
if let Some(entry) = map.get_mut(&user_id) {
|
||||
entry.state = ContainerState::Failed;
|
||||
entry.last_error = Some(e.to_string());
|
||||
}
|
||||
return Err(PoolError::Runtime(e));
|
||||
}
|
||||
}
|
||||
|
||||
let entry_now = self
|
||||
.containers
|
||||
.read()
|
||||
.await
|
||||
.get(&user_id)
|
||||
.cloned()
|
||||
.ok_or_else(|| PoolError::Internal("entry vanished mid-start".into()))?;
|
||||
|
||||
self.probe_until_ready(user_id, entry_now).await
|
||||
}
|
||||
|
||||
/// Wait for the container to answer `/system_stats`, then mark Ready.
|
||||
async fn probe_until_ready(&self, user_id: i64, entry: PoolEntry) -> PoolResult<PoolEntry> {
|
||||
let url = format!("http://127.0.0.1:{}/system_stats", entry.host_port);
|
||||
let deadline = Instant::now() + READINESS_TIMEOUT;
|
||||
|
||||
while Instant::now() < deadline {
|
||||
match self.http.get(&url).send().await {
|
||||
Ok(r) if r.status().is_success() => {
|
||||
let mut map = self.containers.write().await;
|
||||
if let Some(e) = map.get_mut(&user_id) {
|
||||
e.state = ContainerState::Ready;
|
||||
e.last_used_at = Instant::now();
|
||||
e.last_error = None;
|
||||
return Ok(e.clone());
|
||||
}
|
||||
return Err(PoolError::Internal("entry removed during probe".into()));
|
||||
}
|
||||
_ => {
|
||||
tokio::time::sleep(Duration::from_millis(500)).await;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Timeout. Mark Failed and let the next call retry.
|
||||
let mut map = self.containers.write().await;
|
||||
if let Some(e) = map.get_mut(&user_id) {
|
||||
e.state = ContainerState::Failed;
|
||||
e.last_error = Some("readiness probe timed out".into());
|
||||
}
|
||||
// Best-effort stop so the slot frees on the next reaper tick.
|
||||
let name = entry.container_name.clone();
|
||||
let runtime = self.runtime.clone();
|
||||
tokio::spawn(async move {
|
||||
let _ = runtime.stop(&name).await;
|
||||
});
|
||||
|
||||
Err(PoolError::StartTimeout(READINESS_TIMEOUT))
|
||||
}
|
||||
|
||||
/// Stop the user's container, release its port.
|
||||
pub async fn stop(&self, user_id: i64) -> PoolResult<()> {
|
||||
let entry = {
|
||||
let mut map = self.containers.write().await;
|
||||
map.remove(&user_id)
|
||||
};
|
||||
let Some(entry) = entry else {
|
||||
return Ok(());
|
||||
};
|
||||
self.ports.release(entry.host_port);
|
||||
self.runtime.stop(&entry.container_name).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Current state for a user (or `Stopped` if no entry).
|
||||
pub async fn state(&self, user_id: i64) -> ContainerState {
|
||||
self.containers
|
||||
.read()
|
||||
.await
|
||||
.get(&user_id)
|
||||
.map(|e| e.state)
|
||||
.unwrap_or(ContainerState::Stopped)
|
||||
}
|
||||
|
||||
/// Snapshot of the user's pool entry, if any.
|
||||
pub async fn entry(&self, user_id: i64) -> Option<PoolEntry> {
|
||||
self.containers.read().await.get(&user_id).cloned()
|
||||
}
|
||||
|
||||
/// Bump `last_used_at`. Cheap; called on every proxied frame.
|
||||
pub async fn touch(&self, user_id: i64) {
|
||||
self.touch_inner(user_id).await
|
||||
}
|
||||
|
||||
async fn touch_inner(&self, user_id: i64) {
|
||||
if let Some(e) = self.containers.write().await.get_mut(&user_id) {
|
||||
e.last_used_at = Instant::now();
|
||||
}
|
||||
}
|
||||
|
||||
/// Increment / decrement the active-WS counter for a user.
|
||||
pub async fn ws_open(&self, user_id: i64) {
|
||||
if let Some(e) = self.containers.write().await.get_mut(&user_id) {
|
||||
e.active_ws = e.active_ws.saturating_add(1);
|
||||
e.last_used_at = Instant::now();
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn ws_close(&self, user_id: i64) {
|
||||
if let Some(e) = self.containers.write().await.get_mut(&user_id) {
|
||||
e.active_ws = e.active_ws.saturating_sub(1);
|
||||
e.last_used_at = Instant::now();
|
||||
}
|
||||
}
|
||||
|
||||
/// Reconcile against the runtime at startup. Adopts existing
|
||||
/// `comfyui-user-*` containers so a Rust restart doesn't orphan them.
|
||||
pub async fn reconcile_on_startup(&self) {
|
||||
let listings = match self.runtime.list_user_containers().await {
|
||||
Ok(v) => v,
|
||||
Err(e) => {
|
||||
warn!(error = %e, "could not reconcile existing containers; assuming none");
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
let mut map = self.containers.write().await;
|
||||
for listing in listings {
|
||||
// Names look like `comfyui-user-7`.
|
||||
let Some(suffix) = listing.name.strip_prefix("comfyui-user-") else {
|
||||
continue;
|
||||
};
|
||||
let Ok(user_id) = suffix.parse::<i64>() else {
|
||||
continue;
|
||||
};
|
||||
self.ports.claim(listing.host_port);
|
||||
map.insert(
|
||||
user_id,
|
||||
PoolEntry {
|
||||
user_id,
|
||||
container_name: listing.name.clone(),
|
||||
host_port: listing.host_port,
|
||||
state: ContainerState::Ready, // assume ready; first request re-probes
|
||||
last_used_at: Instant::now(),
|
||||
active_ws: 0,
|
||||
last_error: None,
|
||||
},
|
||||
);
|
||||
info!(user_id, container = %listing.name, port = listing.host_port, "adopted existing container");
|
||||
}
|
||||
}
|
||||
|
||||
/// Stop every container we manage (graceful-shutdown hook).
|
||||
pub async fn shutdown_all(&self) {
|
||||
let names: Vec<String> = {
|
||||
let map = self.containers.read().await;
|
||||
map.values().map(|e| e.container_name.clone()).collect()
|
||||
};
|
||||
for name in names {
|
||||
if let Err(e) = self.runtime.stop(&name).await {
|
||||
warn!(container = %name, error = %e, "shutdown stop failed");
|
||||
}
|
||||
}
|
||||
self.containers.write().await.clear();
|
||||
}
|
||||
|
||||
/// Build the [`ContainerSpec`] for a given user.
|
||||
fn build_spec(&self, user_id: i64, name: &str, host_port: u16) -> ContainerSpec {
|
||||
let user_root = self.config.data_root.join(user_id.to_string());
|
||||
|
||||
// Subdirs we mount in. Created on demand by the host operator (or
|
||||
// the runtime can `mkdir -p` them — Docker `bind` will fail if the
|
||||
// path is missing).
|
||||
let mk = |sub: &str| BindMount {
|
||||
host_path: user_root.join(sub),
|
||||
container_path: format!("/workspace/ComfyUI/{}", sub),
|
||||
read_only: false,
|
||||
};
|
||||
|
||||
// Best-effort mkdir so a fresh user just works. Failures here surface
|
||||
// through the Docker error message later.
|
||||
for sub in ["input", "output", "temp", "user", "custom_nodes_user"] {
|
||||
let _ = std::fs::create_dir_all(user_root.join(sub));
|
||||
}
|
||||
|
||||
let mut env = std::collections::HashMap::new();
|
||||
env.insert(
|
||||
"BOOSAPI_BUILTIN_KEY".into(),
|
||||
self.config.builtin_key.clone(),
|
||||
);
|
||||
env.insert(
|
||||
"BOOSAPI_BASE".into(),
|
||||
container_reachable_boosapi_base(&self.config.boosapi_base),
|
||||
);
|
||||
env.insert("BOOSAPI_KEY".into(), self.config.builtin_key.clone());
|
||||
|
||||
ContainerSpec {
|
||||
image: self.config.image.clone(),
|
||||
name: name.to_string(),
|
||||
host_port,
|
||||
env,
|
||||
binds: vec![
|
||||
mk("input"),
|
||||
mk("output"),
|
||||
mk("temp"),
|
||||
mk("user"),
|
||||
mk("custom_nodes_user"),
|
||||
],
|
||||
tmpfs: vec![TmpfsMount {
|
||||
container_path: "/tmp".into(),
|
||||
size_bytes: 512 * 1024 * 1024,
|
||||
}],
|
||||
limits: ResourceLimits::default(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn container_reachable_boosapi_base(base: &str) -> String {
|
||||
if let Ok(override_base) = std::env::var("BOOSAPI_COMFY_CONTAINER_BASE") {
|
||||
let trimmed = override_base.trim();
|
||||
if !trimmed.is_empty() {
|
||||
return trimmed.trim_end_matches('/').to_string();
|
||||
}
|
||||
}
|
||||
|
||||
let trimmed = base.trim().trim_end_matches('/');
|
||||
for local_host in ["127.0.0.1", "localhost"] {
|
||||
let needle = format!("//{}:", local_host);
|
||||
if let Some((prefix, rest)) = trimmed.split_once(&needle) {
|
||||
return format!("{}//host.docker.internal:{}", prefix, rest);
|
||||
}
|
||||
}
|
||||
trimmed.to_string()
|
||||
}
|
||||
|
||||
/// Spawn the idle-reaper task. Lives for the life of the process.
|
||||
pub fn spawn_reaper(pool: Arc<ComfyUIPool>) {
|
||||
tokio::spawn(async move {
|
||||
let mut tick = tokio::time::interval(REAPER_TICK);
|
||||
tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
|
||||
loop {
|
||||
tick.tick().await;
|
||||
reap_once(&pool).await;
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
async fn reap_once(pool: &Arc<ComfyUIPool>) {
|
||||
let now = Instant::now();
|
||||
let idle = pool.config.idle_timeout;
|
||||
let candidates: Vec<i64> = {
|
||||
let map = pool.containers.read().await;
|
||||
map.values()
|
||||
.filter(|e| {
|
||||
e.active_ws == 0
|
||||
&& now.saturating_duration_since(e.last_used_at) > idle
|
||||
&& matches!(e.state, ContainerState::Ready | ContainerState::Failed)
|
||||
})
|
||||
.map(|e| e.user_id)
|
||||
.collect()
|
||||
};
|
||||
|
||||
for user_id in candidates {
|
||||
info!(user_id, "reaping idle comfyui container");
|
||||
if let Err(e) = pool.stop(user_id).await {
|
||||
error!(user_id, error = %e, "reap failed");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use async_trait::async_trait;
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
|
||||
use crate::services::container_runtime::{ContainerHandle, ContainerListing, RuntimeResult};
|
||||
|
||||
/// In-memory runtime that succeeds and pretends `/system_stats` is up.
|
||||
/// Real readiness probe is tested separately via the live Docker setup.
|
||||
struct MockRuntime {
|
||||
starts: AtomicUsize,
|
||||
stops: AtomicUsize,
|
||||
}
|
||||
|
||||
impl MockRuntime {
|
||||
fn new() -> Self {
|
||||
Self {
|
||||
starts: AtomicUsize::new(0),
|
||||
stops: AtomicUsize::new(0),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ContainerRuntime for MockRuntime {
|
||||
async fn start(&self, spec: ContainerSpec) -> RuntimeResult<ContainerHandle> {
|
||||
self.starts.fetch_add(1, Ordering::SeqCst);
|
||||
Ok(ContainerHandle {
|
||||
id: format!("mock-{}", spec.name),
|
||||
name: spec.name,
|
||||
host_port: spec.host_port,
|
||||
})
|
||||
}
|
||||
async fn stop(&self, _name: &str) -> RuntimeResult<()> {
|
||||
self.stops.fetch_add(1, Ordering::SeqCst);
|
||||
Ok(())
|
||||
}
|
||||
async fn is_alive(&self, _name: &str) -> RuntimeResult<bool> {
|
||||
Ok(true)
|
||||
}
|
||||
async fn list_user_containers(&self) -> RuntimeResult<Vec<ContainerListing>> {
|
||||
Ok(vec![])
|
||||
}
|
||||
}
|
||||
|
||||
fn cfg() -> PoolConfig {
|
||||
PoolConfig {
|
||||
image: "test".into(),
|
||||
max_concurrent: 2,
|
||||
idle_timeout: Duration::from_secs(60),
|
||||
data_root: std::env::temp_dir().join("boosapi-pool-test"),
|
||||
boosapi_base: "http://127.0.0.1:9999".into(),
|
||||
builtin_key: "sk-test".into(),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn capacity_cap_enforced() {
|
||||
let runtime = Arc::new(MockRuntime::new());
|
||||
let pool = ComfyUIPool::new(runtime.clone(), cfg());
|
||||
|
||||
// Inject 2 containers manually so we don't have to satisfy the
|
||||
// /system_stats probe (this test isolates the cap check).
|
||||
{
|
||||
let mut map = pool.containers.write().await;
|
||||
map.insert(
|
||||
1,
|
||||
PoolEntry {
|
||||
user_id: 1,
|
||||
container_name: "comfyui-user-1".into(),
|
||||
host_port: 8200,
|
||||
state: ContainerState::Ready,
|
||||
last_used_at: Instant::now(),
|
||||
active_ws: 0,
|
||||
last_error: None,
|
||||
},
|
||||
);
|
||||
map.insert(
|
||||
2,
|
||||
PoolEntry {
|
||||
user_id: 2,
|
||||
container_name: "comfyui-user-2".into(),
|
||||
host_port: 8201,
|
||||
state: ContainerState::Ready,
|
||||
last_used_at: Instant::now(),
|
||||
active_ws: 0,
|
||||
last_error: None,
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
let err = pool.get_or_start(3).await.unwrap_err();
|
||||
assert!(matches!(err, PoolError::AtCapacity(2)), "got: {:?}", err);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn stop_releases_port() {
|
||||
let runtime = Arc::new(MockRuntime::new());
|
||||
let pool = ComfyUIPool::new(runtime.clone(), cfg());
|
||||
{
|
||||
let mut map = pool.containers.write().await;
|
||||
map.insert(
|
||||
7,
|
||||
PoolEntry {
|
||||
user_id: 7,
|
||||
container_name: "comfyui-user-7".into(),
|
||||
host_port: 8205,
|
||||
state: ContainerState::Ready,
|
||||
last_used_at: Instant::now(),
|
||||
active_ws: 0,
|
||||
last_error: None,
|
||||
},
|
||||
);
|
||||
}
|
||||
pool.ports.claim(8205);
|
||||
pool.stop(7).await.unwrap();
|
||||
assert_eq!(runtime.stops.load(Ordering::SeqCst), 1);
|
||||
// Port is back in the free pool.
|
||||
let next = pool.ports.allocate().unwrap();
|
||||
assert_eq!(next, PORT_RANGE_START);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn port_allocator_assigns_in_order_and_releases() {
|
||||
let p = PortAllocator::new();
|
||||
let a = p.allocate().unwrap();
|
||||
let b = p.allocate().unwrap();
|
||||
assert_eq!(a, PORT_RANGE_START);
|
||||
assert_eq!(b, PORT_RANGE_START + 1);
|
||||
p.release(a);
|
||||
let c = p.allocate().unwrap();
|
||||
assert_eq!(c, PORT_RANGE_START);
|
||||
}
|
||||
}
|
||||
@@ -1,120 +0,0 @@
|
||||
//! Backend-agnostic container lifecycle abstraction.
|
||||
//!
|
||||
//! v1 ships [`crate::services::docker_runtime::DockerRuntime`] against the
|
||||
//! local Docker daemon. v2 will add a K3S implementation behind the same
|
||||
//! trait so the pool, proxy, and route handlers do NOT have to change.
|
||||
//!
|
||||
//! Callers (the pool, in particular) MUST go through this trait — never
|
||||
//! import `bollard`, the `docker` CLI, or `kube` types directly.
|
||||
|
||||
use async_trait::async_trait;
|
||||
use std::collections::HashMap;
|
||||
use std::path::PathBuf;
|
||||
|
||||
/// Resource and security caps for a container.
|
||||
///
|
||||
/// Mirrors §design Decision 5 verbatim: 2 CPU, 2 GB RAM, 2 GB swap,
|
||||
/// `pids_limit=512`, read-only rootfs, drop all caps, `no-new-privileges`.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ResourceLimits {
|
||||
pub cpus: f64,
|
||||
pub memory_bytes: u64,
|
||||
pub memory_swap_bytes: u64,
|
||||
pub pids_limit: u32,
|
||||
pub read_only_rootfs: bool,
|
||||
pub cap_drop_all: bool,
|
||||
pub no_new_privileges: bool,
|
||||
}
|
||||
|
||||
impl Default for ResourceLimits {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
cpus: 2.0,
|
||||
memory_bytes: 2 * 1024 * 1024 * 1024,
|
||||
memory_swap_bytes: 2 * 1024 * 1024 * 1024,
|
||||
pids_limit: 512,
|
||||
read_only_rootfs: true,
|
||||
cap_drop_all: true,
|
||||
no_new_privileges: true,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// One bind mount: `host_path -> container_path`.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct BindMount {
|
||||
pub host_path: PathBuf,
|
||||
pub container_path: String,
|
||||
pub read_only: bool,
|
||||
}
|
||||
|
||||
/// Tmpfs mount (used to satisfy a few writable paths under read-only rootfs).
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct TmpfsMount {
|
||||
pub container_path: String,
|
||||
pub size_bytes: u64,
|
||||
}
|
||||
|
||||
/// A complete container start request.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ContainerSpec {
|
||||
pub image: String,
|
||||
pub name: String,
|
||||
/// Host port to bind to. The container's internal port is fixed at 8188.
|
||||
pub host_port: u16,
|
||||
pub env: HashMap<String, String>,
|
||||
pub binds: Vec<BindMount>,
|
||||
pub tmpfs: Vec<TmpfsMount>,
|
||||
pub limits: ResourceLimits,
|
||||
}
|
||||
|
||||
/// Result of a successful start: identifies the container and tells the
|
||||
/// caller which loopback port to forward traffic to.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ContainerHandle {
|
||||
/// Backend-specific id (Docker container id, k8s pod name, ...).
|
||||
pub id: String,
|
||||
/// Container name we asked for, e.g. `comfyui-user-7`.
|
||||
pub name: String,
|
||||
/// Loopback host port the runtime mapped for us.
|
||||
pub host_port: u16,
|
||||
}
|
||||
|
||||
/// Listing entry returned by [`ContainerRuntime::list_user_containers`].
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ContainerListing {
|
||||
pub id: String,
|
||||
pub name: String,
|
||||
pub host_port: u16,
|
||||
}
|
||||
|
||||
#[derive(thiserror::Error, Debug)]
|
||||
pub enum RuntimeError {
|
||||
#[error("backend error: {0}")]
|
||||
Backend(String),
|
||||
#[error("container not found: {0}")]
|
||||
NotFound(String),
|
||||
}
|
||||
|
||||
pub type RuntimeResult<T> = Result<T, RuntimeError>;
|
||||
|
||||
/// Backend-agnostic container lifecycle.
|
||||
#[async_trait]
|
||||
pub trait ContainerRuntime: Send + Sync + 'static {
|
||||
/// Start a container and return a handle once it is running.
|
||||
///
|
||||
/// The runtime is NOT required to wait for the application inside the
|
||||
/// container to be ready — that probe lives in the pool.
|
||||
async fn start(&self, spec: ContainerSpec) -> RuntimeResult<ContainerHandle>;
|
||||
|
||||
/// Stop and remove a container by name. Idempotent: removing a
|
||||
/// nonexistent container returns `Ok`.
|
||||
async fn stop(&self, name: &str) -> RuntimeResult<()>;
|
||||
|
||||
/// Check whether a container is running.
|
||||
async fn is_alive(&self, name: &str) -> RuntimeResult<bool>;
|
||||
|
||||
/// List every running container whose name starts with `comfyui-user-`.
|
||||
/// Used at API startup to reconstruct the pool's in-memory map.
|
||||
async fn list_user_containers(&self) -> RuntimeResult<Vec<ContainerListing>>;
|
||||
}
|
||||
@@ -1,291 +0,0 @@
|
||||
//! Docker implementation of [`ContainerRuntime`].
|
||||
//!
|
||||
//! We shell out to the `docker` CLI rather than pull in a Rust Docker SDK,
|
||||
//! to keep the dependency footprint small. The CLI is universally
|
||||
//! available on the hosts we ship to, and the surface area we need
|
||||
//! (`run / stop / rm / inspect / ps`) is tiny.
|
||||
|
||||
use async_trait::async_trait;
|
||||
use std::process::Stdio;
|
||||
use tokio::process::Command;
|
||||
use tracing::{debug, warn};
|
||||
|
||||
use super::container_runtime::{
|
||||
ContainerHandle, ContainerListing, ContainerRuntime, ContainerSpec, RuntimeError, RuntimeResult,
|
||||
};
|
||||
|
||||
/// Internal port the ComfyUI process listens on inside the container.
|
||||
/// Hard-coded to match the image's ENTRYPOINT.
|
||||
const CONTAINER_INTERNAL_PORT: u16 = 8188;
|
||||
|
||||
/// Label we stamp onto every container we manage. Used by
|
||||
/// `list_user_containers` and as a sanity-check on `stop`.
|
||||
pub const MANAGED_LABEL: &str = "boosapi.comfy.tenant=1";
|
||||
|
||||
/// Default Docker runtime. Stateless — every method shells out fresh.
|
||||
pub struct DockerRuntime {
|
||||
docker_bin: String,
|
||||
}
|
||||
|
||||
impl DockerRuntime {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
docker_bin: std::env::var("BOOSAPI_DOCKER_BIN").unwrap_or_else(|_| "docker".into()),
|
||||
}
|
||||
}
|
||||
|
||||
fn cmd(&self) -> Command {
|
||||
let mut c = Command::new(&self.docker_bin);
|
||||
c.stdin(Stdio::null());
|
||||
c
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for DockerRuntime {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl ContainerRuntime for DockerRuntime {
|
||||
async fn start(&self, spec: ContainerSpec) -> RuntimeResult<ContainerHandle> {
|
||||
// Best-effort: kill any stale container with the same name.
|
||||
// `docker rm -f <name>` is a no-op if the name is free.
|
||||
let _ = self
|
||||
.cmd()
|
||||
.args(["rm", "-f", &spec.name])
|
||||
.stdout(Stdio::null())
|
||||
.stderr(Stdio::null())
|
||||
.status()
|
||||
.await;
|
||||
|
||||
let mut cmd = self.cmd();
|
||||
cmd.arg("run")
|
||||
.arg("-d")
|
||||
.arg("--rm")
|
||||
.arg("--name")
|
||||
.arg(&spec.name)
|
||||
.arg("--label")
|
||||
.arg(MANAGED_LABEL)
|
||||
.arg("-p")
|
||||
.arg(format!(
|
||||
"127.0.0.1:{}:{}",
|
||||
spec.host_port, CONTAINER_INTERNAL_PORT
|
||||
))
|
||||
.arg("--add-host")
|
||||
.arg("host.docker.internal:host-gateway")
|
||||
// Resource limits (§design Decision 5)
|
||||
.arg(format!("--cpus={}", spec.limits.cpus))
|
||||
.arg(format!("--memory={}", spec.limits.memory_bytes))
|
||||
.arg(format!("--memory-swap={}", spec.limits.memory_swap_bytes))
|
||||
.arg(format!("--pids-limit={}", spec.limits.pids_limit));
|
||||
|
||||
if spec.limits.read_only_rootfs {
|
||||
cmd.arg("--read-only");
|
||||
}
|
||||
if spec.limits.cap_drop_all {
|
||||
cmd.args(["--cap-drop", "ALL"]);
|
||||
}
|
||||
if spec.limits.no_new_privileges {
|
||||
cmd.args(["--security-opt", "no-new-privileges"]);
|
||||
}
|
||||
|
||||
for (k, v) in &spec.env {
|
||||
cmd.arg("-e").arg(format!("{}={}", k, v));
|
||||
}
|
||||
|
||||
for bind in &spec.binds {
|
||||
// Using `--mount` over `-v` so the syntax is unambiguous and
|
||||
// we get a clear error if the host path is wrong.
|
||||
let mut s = format!(
|
||||
"type=bind,source={},target={}",
|
||||
bind.host_path.display(),
|
||||
bind.container_path
|
||||
);
|
||||
if bind.read_only {
|
||||
s.push_str(",readonly");
|
||||
}
|
||||
cmd.arg("--mount").arg(s);
|
||||
}
|
||||
|
||||
for tmp in &spec.tmpfs {
|
||||
cmd.arg("--tmpfs")
|
||||
.arg(format!("{}:size={}", tmp.container_path, tmp.size_bytes));
|
||||
}
|
||||
|
||||
cmd.arg(&spec.image);
|
||||
|
||||
debug!(?cmd, "starting comfyui container");
|
||||
|
||||
let output = cmd.output().await.map_err(io_to_runtime)?;
|
||||
if !output.status.success() {
|
||||
return Err(RuntimeError::Backend(format!(
|
||||
"docker run failed (exit {}): {}",
|
||||
output.status,
|
||||
String::from_utf8_lossy(&output.stderr).trim()
|
||||
)));
|
||||
}
|
||||
|
||||
let id = String::from_utf8_lossy(&output.stdout).trim().to_string();
|
||||
if id.is_empty() {
|
||||
return Err(RuntimeError::Backend(
|
||||
"docker run returned empty container id".into(),
|
||||
));
|
||||
}
|
||||
|
||||
Ok(ContainerHandle {
|
||||
id,
|
||||
name: spec.name,
|
||||
host_port: spec.host_port,
|
||||
})
|
||||
}
|
||||
|
||||
async fn stop(&self, name: &str) -> RuntimeResult<()> {
|
||||
// `docker stop` flushes nicely; `--rm` from start removes it after.
|
||||
// If it's already gone, treat as success.
|
||||
let output = self
|
||||
.cmd()
|
||||
.args(["stop", "-t", "10", name])
|
||||
.output()
|
||||
.await
|
||||
.map_err(io_to_runtime)?;
|
||||
|
||||
if !output.status.success() {
|
||||
let stderr = String::from_utf8_lossy(&output.stderr);
|
||||
if stderr.contains("No such container") {
|
||||
return Ok(());
|
||||
}
|
||||
warn!(name, stderr = stderr.trim(), "docker stop failed");
|
||||
// Fall through to rm -f as a last resort.
|
||||
}
|
||||
|
||||
let _ = self
|
||||
.cmd()
|
||||
.args(["rm", "-f", name])
|
||||
.stdout(Stdio::null())
|
||||
.stderr(Stdio::null())
|
||||
.status()
|
||||
.await;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn is_alive(&self, name: &str) -> RuntimeResult<bool> {
|
||||
let output = self
|
||||
.cmd()
|
||||
.args(["inspect", "-f", "{{.State.Running}}", name])
|
||||
.output()
|
||||
.await
|
||||
.map_err(io_to_runtime)?;
|
||||
|
||||
if !output.status.success() {
|
||||
// `docker inspect` returns 1 if the container is unknown.
|
||||
return Ok(false);
|
||||
}
|
||||
|
||||
Ok(String::from_utf8_lossy(&output.stdout).trim() == "true")
|
||||
}
|
||||
|
||||
async fn list_user_containers(&self) -> RuntimeResult<Vec<ContainerListing>> {
|
||||
// `--format {{json .}}` gives one JSON object per line. We pull the
|
||||
// fields we need without depending on a JSON-deserialize crate
|
||||
// schema for Docker's somewhat-fluid output.
|
||||
let output = self
|
||||
.cmd()
|
||||
.args([
|
||||
"ps",
|
||||
"--filter",
|
||||
"name=comfyui-user-",
|
||||
"--filter",
|
||||
&format!("label={}", MANAGED_LABEL),
|
||||
"--format",
|
||||
"{{.Names}}\t{{.ID}}\t{{.Ports}}",
|
||||
])
|
||||
.output()
|
||||
.await
|
||||
.map_err(io_to_runtime)?;
|
||||
|
||||
if !output.status.success() {
|
||||
return Err(RuntimeError::Backend(format!(
|
||||
"docker ps failed: {}",
|
||||
String::from_utf8_lossy(&output.stderr).trim()
|
||||
)));
|
||||
}
|
||||
|
||||
let stdout = String::from_utf8_lossy(&output.stdout);
|
||||
let mut out = Vec::new();
|
||||
for line in stdout.lines() {
|
||||
let mut parts = line.splitn(3, '\t');
|
||||
let name = parts.next().unwrap_or("").trim().to_string();
|
||||
let id = parts.next().unwrap_or("").trim().to_string();
|
||||
let ports = parts.next().unwrap_or("").trim();
|
||||
if name.is_empty() || id.is_empty() {
|
||||
continue;
|
||||
}
|
||||
let host_port = match parse_host_port(ports) {
|
||||
Some(p) => p,
|
||||
None => {
|
||||
warn!(
|
||||
name,
|
||||
ports, "could not parse host port from docker ps; skipping"
|
||||
);
|
||||
continue;
|
||||
}
|
||||
};
|
||||
out.push(ContainerListing {
|
||||
id,
|
||||
name,
|
||||
host_port,
|
||||
});
|
||||
}
|
||||
Ok(out)
|
||||
}
|
||||
}
|
||||
|
||||
/// Pull the host-side port out of a `docker ps` Ports column.
|
||||
///
|
||||
/// Examples it must handle:
|
||||
/// "127.0.0.1:8200->8188/tcp"
|
||||
/// "127.0.0.1:8200->8188/tcp, 127.0.0.1:8200->8188/udp"
|
||||
fn parse_host_port(ports: &str) -> Option<u16> {
|
||||
for spec in ports.split(',') {
|
||||
let spec = spec.trim();
|
||||
// Accept `<host>:<port>->...`
|
||||
let arrow = spec.find("->")?;
|
||||
let left = &spec[..arrow];
|
||||
let port_str = left.rsplit(':').next()?;
|
||||
if let Ok(p) = port_str.parse::<u16>() {
|
||||
return Some(p);
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
fn io_to_runtime(e: std::io::Error) -> RuntimeError {
|
||||
RuntimeError::Backend(format!("docker cli io error: {}", e))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn parses_simple_port_mapping() {
|
||||
assert_eq!(parse_host_port("127.0.0.1:8200->8188/tcp"), Some(8200));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parses_port_mapping_with_extra_entries() {
|
||||
assert_eq!(
|
||||
parse_host_port("127.0.0.1:8201->8188/tcp, 127.0.0.1:8201->8188/udp"),
|
||||
Some(8201)
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn returns_none_for_garbage() {
|
||||
assert_eq!(parse_host_port(""), None);
|
||||
assert_eq!(parse_host_port("nope"), None);
|
||||
}
|
||||
}
|
||||
+1
-3
@@ -1,13 +1,11 @@
|
||||
pub mod announcement;
|
||||
pub mod billing;
|
||||
pub mod builtin_keys;
|
||||
pub mod comfyui_pool;
|
||||
pub mod comfyui_process;
|
||||
pub mod container_runtime;
|
||||
pub mod cron;
|
||||
pub mod docker_runtime;
|
||||
pub mod notification;
|
||||
pub mod oauth;
|
||||
pub mod registration_limiter;
|
||||
pub mod route_rebuild;
|
||||
pub mod shared;
|
||||
pub mod update_center;
|
||||
|
||||
@@ -0,0 +1,221 @@
|
||||
use std::{
|
||||
collections::{HashMap, VecDeque},
|
||||
net::IpAddr,
|
||||
time::{Duration, Instant},
|
||||
};
|
||||
|
||||
use parking_lot::Mutex;
|
||||
|
||||
const IP_WINDOW: Duration = Duration::from_secs(15 * 60);
|
||||
const EMAIL_WINDOW: Duration = Duration::from_secs(60 * 60);
|
||||
const GLOBAL_WINDOW: Duration = Duration::from_secs(5 * 60);
|
||||
const SUCCESS_LOCK_WINDOW: Duration = Duration::from_secs(10 * 60);
|
||||
|
||||
const MAX_ATTEMPTS_PER_IP: usize = 8;
|
||||
const MAX_ATTEMPTS_PER_EMAIL: usize = 3;
|
||||
const MAX_GLOBAL_ATTEMPTS: usize = 120;
|
||||
const MAX_SUCCESS_PER_IP: usize = 2;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub struct RegistrationLimitExceeded {
|
||||
pub retry_after_secs: u64,
|
||||
pub message: &'static str,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct LimiterState {
|
||||
attempts_by_ip: HashMap<String, VecDeque<Instant>>,
|
||||
attempts_by_email: HashMap<String, VecDeque<Instant>>,
|
||||
successful_by_ip: HashMap<String, VecDeque<Instant>>,
|
||||
global_attempts: VecDeque<Instant>,
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
pub struct RegistrationLimiter {
|
||||
inner: Mutex<LimiterState>,
|
||||
}
|
||||
|
||||
impl RegistrationLimiter {
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
pub fn check_and_record_attempt(
|
||||
&self,
|
||||
ip: Option<IpAddr>,
|
||||
email: &str,
|
||||
) -> Result<(), RegistrationLimitExceeded> {
|
||||
let now = Instant::now();
|
||||
let ip_key = ip
|
||||
.map(|addr| addr.to_string())
|
||||
.unwrap_or_else(|| "unknown".to_string());
|
||||
let email_key = normalize_email(email);
|
||||
|
||||
let mut guard = self.inner.lock();
|
||||
prune_deque(&mut guard.global_attempts, now, GLOBAL_WINDOW);
|
||||
if guard.global_attempts.len() >= MAX_GLOBAL_ATTEMPTS {
|
||||
return Err(limit_error(
|
||||
&guard.global_attempts,
|
||||
now,
|
||||
GLOBAL_WINDOW,
|
||||
"注册请求过于频繁，请稍后再试。",
|
||||
));
|
||||
}
|
||||
|
||||
{
|
||||
let ip_attempts = guard.attempts_by_ip.entry(ip_key.clone()).or_default();
|
||||
prune_deque(ip_attempts, now, IP_WINDOW);
|
||||
if ip_attempts.len() >= MAX_ATTEMPTS_PER_IP {
|
||||
return Err(limit_error(
|
||||
ip_attempts,
|
||||
now,
|
||||
IP_WINDOW,
|
||||
"当前网络注册过于频繁，请稍后再试。",
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
{
|
||||
let ip_successes = guard.successful_by_ip.entry(ip_key.clone()).or_default();
|
||||
prune_deque(ip_successes, now, SUCCESS_LOCK_WINDOW);
|
||||
if ip_successes.len() >= MAX_SUCCESS_PER_IP {
|
||||
return Err(limit_error(
|
||||
ip_successes,
|
||||
now,
|
||||
SUCCESS_LOCK_WINDOW,
|
||||
"当前网络短时间内已注册多个账号，请稍后再试。",
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
{
|
||||
let email_attempts = guard
|
||||
.attempts_by_email
|
||||
.entry(email_key.clone())
|
||||
.or_default();
|
||||
prune_deque(email_attempts, now, EMAIL_WINDOW);
|
||||
if email_attempts.len() >= MAX_ATTEMPTS_PER_EMAIL {
|
||||
return Err(limit_error(
|
||||
email_attempts,
|
||||
now,
|
||||
EMAIL_WINDOW,
|
||||
"该邮箱注册尝试过于频繁，请稍后再试。",
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
guard.global_attempts.push_back(now);
|
||||
guard
|
||||
.attempts_by_ip
|
||||
.entry(ip_key)
|
||||
.or_default()
|
||||
.push_back(now);
|
||||
guard
|
||||
.attempts_by_email
|
||||
.entry(email_key)
|
||||
.or_default()
|
||||
.push_back(now);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn record_success(&self, ip: Option<IpAddr>) {
|
||||
let now = Instant::now();
|
||||
let ip_key = ip
|
||||
.map(|addr| addr.to_string())
|
||||
.unwrap_or_else(|| "unknown".to_string());
|
||||
let mut guard = self.inner.lock();
|
||||
let entries = guard.successful_by_ip.entry(ip_key).or_default();
|
||||
prune_deque(entries, now, SUCCESS_LOCK_WINDOW);
|
||||
entries.push_back(now);
|
||||
}
|
||||
}
|
||||
|
||||
fn normalize_email(email: &str) -> String {
|
||||
email.trim().to_ascii_lowercase()
|
||||
}
|
||||
|
||||
fn prune_deque(entries: &mut VecDeque<Instant>, now: Instant, window: Duration) {
|
||||
while entries
|
||||
.front()
|
||||
.is_some_and(|timestamp| now.duration_since(*timestamp) >= window)
|
||||
{
|
||||
entries.pop_front();
|
||||
}
|
||||
}
|
||||
|
||||
fn limit_error(
|
||||
entries: &VecDeque<Instant>,
|
||||
now: Instant,
|
||||
window: Duration,
|
||||
message: &'static str,
|
||||
) -> RegistrationLimitExceeded {
|
||||
let retry_after_secs = entries
|
||||
.front()
|
||||
.map(|first| {
|
||||
window
|
||||
.saturating_sub(now.duration_since(*first))
|
||||
.as_secs()
|
||||
.max(1)
|
||||
})
|
||||
.unwrap_or_else(|| window.as_secs().max(1));
|
||||
RegistrationLimitExceeded {
|
||||
retry_after_secs,
|
||||
message,
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::net::{IpAddr, Ipv4Addr};
|
||||
|
||||
fn ip() -> Option<IpAddr> {
|
||||
Some(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)))
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn limits_repeated_attempts_from_same_ip() {
|
||||
let limiter = RegistrationLimiter::new();
|
||||
for index in 0..MAX_ATTEMPTS_PER_IP {
|
||||
limiter
|
||||
.check_and_record_attempt(ip(), &format!("u{index}@example.com"))
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
let err = limiter
|
||||
.check_and_record_attempt(ip(), "blocked@example.com")
|
||||
.unwrap_err();
|
||||
assert_eq!(err.message, "当前网络注册过于频繁，请稍后再试。");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn limits_repeated_attempts_for_same_email_case_insensitive() {
|
||||
let limiter = RegistrationLimiter::new();
|
||||
for _ in 0..MAX_ATTEMPTS_PER_EMAIL {
|
||||
limiter
|
||||
.check_and_record_attempt(None, "User@Example.com")
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
let err = limiter
|
||||
.check_and_record_attempt(None, "user@example.com")
|
||||
.unwrap_err();
|
||||
assert_eq!(err.message, "该邮箱注册尝试过于频繁，请稍后再试。");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn limits_successful_registrations_from_same_ip() {
|
||||
let limiter = RegistrationLimiter::new();
|
||||
for index in 0..MAX_SUCCESS_PER_IP {
|
||||
limiter
|
||||
.check_and_record_attempt(ip(), &format!("u{index}@example.com"))
|
||||
.unwrap();
|
||||
limiter.record_success(ip());
|
||||
}
|
||||
|
||||
let err = limiter
|
||||
.check_and_record_attempt(ip(), "extra@example.com")
|
||||
.unwrap_err();
|
||||
assert_eq!(err.message, "当前网络短时间内已注册多个账号，请稍后再试。");
|
||||
}
|
||||
}
|
||||
+30
-14
@@ -22,6 +22,7 @@ pub struct WorkspaceAssetInput {
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct StoredWorkspaceAsset {
|
||||
pub public_id: String,
|
||||
pub file_name: Option<String>,
|
||||
pub file_path: Option<String>,
|
||||
pub remote_url: Option<String>,
|
||||
pub status: String,
|
||||
@@ -56,7 +57,12 @@ fn safe_segment(value: &str) -> String {
|
||||
}
|
||||
|
||||
fn extension_for_mime(mime: &str) -> &'static str {
|
||||
let mime = mime.split(';').next().unwrap_or("").trim().to_ascii_lowercase();
|
||||
let mime = mime
|
||||
.split(';')
|
||||
.next()
|
||||
.unwrap_or("")
|
||||
.trim()
|
||||
.to_ascii_lowercase();
|
||||
match mime.as_str() {
|
||||
"image/jpeg" | "image/jpg" => "jpg",
|
||||
"image/webp" => "webp",
|
||||
@@ -156,6 +162,7 @@ async fn insert_asset(
|
||||
|
||||
Ok(StoredWorkspaceAsset {
|
||||
public_id: public_id.to_string(),
|
||||
file_name,
|
||||
file_path: file_path_string,
|
||||
remote_url: input.remote_url.clone(),
|
||||
status: status.to_string(),
|
||||
@@ -193,12 +200,13 @@ pub async fn record_remote_asset(
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.map(str::to_string)
|
||||
.unwrap_or_else(|| "application/octet-stream".to_string());
|
||||
let bytes = response
|
||||
.bytes()
|
||||
.await
|
||||
.map_err(|err| AppError::UpstreamError(format!("download workspace asset: {err}")))?;
|
||||
let ext = extension_from_url(url).unwrap_or_else(|| extension_for_mime(&mime_type).to_string());
|
||||
let path = write_asset_file(input.user_id, &input.source, &public_id, &ext, &bytes).await?;
|
||||
let bytes = response.bytes().await.map_err(|err| {
|
||||
AppError::UpstreamError(format!("download workspace asset: {err}"))
|
||||
})?;
|
||||
let ext = extension_from_url(url)
|
||||
.unwrap_or_else(|| extension_for_mime(&mime_type).to_string());
|
||||
let path =
|
||||
write_asset_file(input.user_id, &input.source, &public_id, &ext, &bytes).await?;
|
||||
insert_asset(
|
||||
state,
|
||||
&input,
|
||||
@@ -283,17 +291,22 @@ pub async fn record_image_assets_from_response(
|
||||
prompt: Option<String>,
|
||||
model: Option<String>,
|
||||
response: &Value,
|
||||
) {
|
||||
) -> Vec<StoredWorkspaceAsset> {
|
||||
let Some(items) = response.get("data").and_then(Value::as_array) else {
|
||||
return;
|
||||
return Vec::new();
|
||||
};
|
||||
let mut stored_assets = Vec::new();
|
||||
for (index, item) in items.iter().enumerate() {
|
||||
let title = prompt
|
||||
.as_deref()
|
||||
.map(|value| value.chars().take(36).collect::<String>())
|
||||
.filter(|value| !value.trim().is_empty())
|
||||
.or_else(|| Some(format!("图片产出 #{}", index + 1)));
|
||||
if let Some(url) = item.get("url").and_then(Value::as_str).filter(|value| !value.trim().is_empty()) {
|
||||
if let Some(url) = item
|
||||
.get("url")
|
||||
.and_then(Value::as_str)
|
||||
.filter(|value| !value.trim().is_empty())
|
||||
{
|
||||
let input = WorkspaceAssetInput {
|
||||
user_id,
|
||||
source: source.to_string(),
|
||||
@@ -304,8 +317,9 @@ pub async fn record_image_assets_from_response(
|
||||
remote_url: Some(url.trim().to_string()),
|
||||
metadata: json!({ "responseItem": item }),
|
||||
};
|
||||
if let Err(err) = record_remote_asset(state, input).await {
|
||||
tracing::warn!(%err, "failed to archive remote image asset");
|
||||
match record_remote_asset(state, input).await {
|
||||
Ok(asset) => stored_assets.push(asset),
|
||||
Err(err) => tracing::warn!(%err, "failed to archive remote image asset"),
|
||||
}
|
||||
} else if let Some(b64) = item
|
||||
.get("b64_json")
|
||||
@@ -322,9 +336,11 @@ pub async fn record_image_assets_from_response(
|
||||
remote_url: None,
|
||||
metadata: json!({ "responseItem": item }),
|
||||
};
|
||||
if let Err(err) = record_base64_asset(state, input, b64).await {
|
||||
tracing::warn!(%err, "failed to archive base64 image asset");
|
||||
match record_base64_asset(state, input, b64).await {
|
||||
Ok(asset) => stored_assets.push(asset),
|
||||
Err(err) => tracing::warn!(%err, "failed to archive base64 image asset"),
|
||||
}
|
||||
}
|
||||
}
|
||||
stored_assets
|
||||
}
|
||||
|
||||
@@ -294,7 +294,13 @@ fn parse_message_content(
|
||||
return Ok((Vec::new(), None, None));
|
||||
};
|
||||
if let Some(text) = content.as_str() {
|
||||
return Ok((vec![ContentPart::Text { text: text.to_string() }], None, None));
|
||||
return Ok((
|
||||
vec![ContentPart::Text {
|
||||
text: text.to_string(),
|
||||
}],
|
||||
None,
|
||||
None,
|
||||
));
|
||||
}
|
||||
|
||||
let mut parts = Vec::new();
|
||||
@@ -307,7 +313,9 @@ fn parse_message_content(
|
||||
match item["type"].as_str() {
|
||||
Some("text") => {
|
||||
if let Some(text) = item["text"].as_str() {
|
||||
parts.push(ContentPart::Text { text: text.to_string() });
|
||||
parts.push(ContentPart::Text {
|
||||
text: text.to_string(),
|
||||
});
|
||||
}
|
||||
}
|
||||
Some("image") => {
|
||||
@@ -317,9 +325,7 @@ fn parse_message_content(
|
||||
detail: None,
|
||||
});
|
||||
} else if item["source"]["type"].as_str() == Some("base64") {
|
||||
let media_type = item["source"]["media_type"]
|
||||
.as_str()
|
||||
.unwrap_or("image/png");
|
||||
let media_type = item["source"]["media_type"].as_str().unwrap_or("image/png");
|
||||
if let Some(data) = item["source"]["data"].as_str() {
|
||||
parts.push(ContentPart::Image {
|
||||
url: format!("data:{media_type};base64,{data}"),
|
||||
@@ -350,14 +356,20 @@ fn parse_message_content(
|
||||
}
|
||||
Ok((
|
||||
parts,
|
||||
if tool_calls.is_empty() { None } else { Some(tool_calls) },
|
||||
if tool_calls.is_empty() {
|
||||
None
|
||||
} else {
|
||||
Some(tool_calls)
|
||||
},
|
||||
tool_call_id,
|
||||
))
|
||||
}
|
||||
|
||||
fn parse_anthropic_content(value: &Value) -> Result<Vec<ContentPart>, String> {
|
||||
if let Some(text) = value.as_str() {
|
||||
return Ok(vec![ContentPart::Text { text: text.to_string() }]);
|
||||
return Ok(vec![ContentPart::Text {
|
||||
text: text.to_string(),
|
||||
}]);
|
||||
}
|
||||
let Some(items) = value.as_array() else {
|
||||
return Ok(Vec::new());
|
||||
@@ -366,7 +378,9 @@ fn parse_anthropic_content(value: &Value) -> Result<Vec<ContentPart>, String> {
|
||||
for item in items {
|
||||
if item["type"].as_str() == Some("text") {
|
||||
if let Some(text) = item["text"].as_str() {
|
||||
parts.push(ContentPart::Text { text: text.to_string() });
|
||||
parts.push(ContentPart::Text {
|
||||
text: text.to_string(),
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -387,7 +401,9 @@ fn parse_anthropic_tool_choice(value: Option<&Value>) -> Option<ToolChoice> {
|
||||
match value["type"].as_str()? {
|
||||
"auto" => Some(ToolChoice::Auto),
|
||||
"any" => Some(ToolChoice::Required),
|
||||
"tool" => value["name"].as_str().map(|name| ToolChoice::Specific(name.to_string())),
|
||||
"tool" => value["name"]
|
||||
.as_str()
|
||||
.map(|name| ToolChoice::Specific(name.to_string())),
|
||||
"none" => Some(ToolChoice::None),
|
||||
_ => None,
|
||||
}
|
||||
|
||||
@@ -5,6 +5,8 @@ pub mod outbound;
|
||||
pub mod stream;
|
||||
|
||||
pub use compatibility::AnthropicMessagesCompatibility;
|
||||
pub use inbound::{format_messages_request, format_messages_request_with_compat, parse_messages_request};
|
||||
pub use inbound::{
|
||||
format_messages_request, format_messages_request_with_compat, parse_messages_request,
|
||||
};
|
||||
pub use outbound::parse_messages_response;
|
||||
pub use stream::parse_stream_event;
|
||||
|
||||
@@ -2,7 +2,6 @@ import { describe, expect, it, vi } from 'vitest';
|
||||
import { create, type ReactTestInstance } from 'react-test-renderer';
|
||||
import { MemoryRouter } from 'react-router-dom';
|
||||
import AdminLoginPage from './pages/auth/AdminLoginPage.js';
|
||||
import { SITE_DOCS_URL, SITE_GITHUB_URL } from './docsLink.js';
|
||||
|
||||
function collectText(node: ReactTestInstance): string {
|
||||
return (node.children || []).map((child) => {
|
||||
@@ -12,14 +11,6 @@ function collectText(node: ReactTestInstance): string {
|
||||
}
|
||||
|
||||
describe('Login surface', () => {
|
||||
it('uses the site root as the documentation URL', () => {
|
||||
expect(SITE_DOCS_URL).toBe('https://boosapi.cita777.me');
|
||||
});
|
||||
|
||||
it('uses the author github profile for the login github shortcut', () => {
|
||||
expect(SITE_GITHUB_URL).toBe('https://github.com/cita-777');
|
||||
});
|
||||
|
||||
it('renders a poster-style hero with a floating admin login panel', () => {
|
||||
const root = create(
|
||||
<MemoryRouter>
|
||||
@@ -52,15 +43,9 @@ describe('Login surface', () => {
|
||||
expect(pageText).toContain('统一代理网关');
|
||||
expect(pageText).toContain('智能路由引擎');
|
||||
expect(pageText).toContain('自动模型发现');
|
||||
expect(pageText).toContain('部署文档');
|
||||
expect(lightBrandPanel).toBeTruthy();
|
||||
expect(authStage).toBeTruthy();
|
||||
expect(brandMarkCanvas).toBeTruthy();
|
||||
|
||||
const docsLink = root.root.find((node) => (
|
||||
node.type === 'a'
|
||||
&& node.props.href === SITE_DOCS_URL
|
||||
));
|
||||
const tokenInput = root.root.find((node) => (
|
||||
node.type === 'input'
|
||||
&& node.props.placeholder === '管理员令牌'
|
||||
@@ -70,14 +55,6 @@ describe('Login surface', () => {
|
||||
&& typeof node.props.className === 'string'
|
||||
&& node.props.className.includes('login-mode-tab')
|
||||
));
|
||||
const githubLink = root.root.find((node) => (
|
||||
node.type === 'a'
|
||||
&& node.props.href === SITE_GITHUB_URL
|
||||
));
|
||||
|
||||
expect(docsLink.props.target).toBe('_blank');
|
||||
expect(githubLink.props['aria-label']).toBe('GitHub');
|
||||
expect(githubLink.props.target).toBe('_blank');
|
||||
expect(tokenInput.props.type).toBe('password');
|
||||
expect(loginModeTabs.map((node) => collectText(node))).toEqual(['Auth Token', '账号密码']);
|
||||
} finally {
|
||||
|
||||
+17
-15
@@ -7,13 +7,11 @@ import TooltipLayer from './components/TooltipLayer.js';
|
||||
import { api, userApi } from './api.js';
|
||||
import { clearSession, getSessionToken } from './authSession.js';
|
||||
import {
|
||||
FIRST_USE_DOC_REMINDER_KEY,
|
||||
LEGACY_THEME_STORAGE_KEY,
|
||||
THEME_MODE_STORAGE_KEY,
|
||||
USER_PROFILE_STORAGE_KEY,
|
||||
} from './appLocalState.js';
|
||||
import { I18nProvider, useI18n } from './i18n.js';
|
||||
import { SITE_DOCS_URL } from './docsLink.js';
|
||||
import { useAnimatedVisibility } from './components/useAnimatedVisibility.js';
|
||||
import { useIsMobile } from './components/useIsMobile.js';
|
||||
import { MobileDrawer } from './components/MobileDrawer.js';
|
||||
@@ -36,10 +34,11 @@ const ProgramLogs = lazy(() => import('./pages/ProgramLogs.js'));
|
||||
const Models = lazy(() => import('./pages/Models.js'));
|
||||
const About = lazy(() => import('./pages/About.js'));
|
||||
const ModelTester = lazy(() => import('./pages/ModelTester.js'));
|
||||
const ComfyUI = lazy(() => import('./pages/ComfyUI.js'));
|
||||
const ComfyUIAgent = lazy(() => import('./pages/ComfyUIAgent.js'));
|
||||
const PictureAgent = lazy(() => import('./pages/PictureAgent.js'));
|
||||
const DigitalHuman = lazy(() => import('./pages/DigitalHuman.js'));
|
||||
const Workspace = lazy(() => import('./pages/Workspace.js'));
|
||||
const Workflows = lazy(() => import('./pages/Workflows.js'));
|
||||
const Monitors = lazy(() => import('./pages/Monitors.js'));
|
||||
const OAuthManagement = lazy(() => import('./pages/OAuthManagement.js'));
|
||||
const SiteAnnouncements = lazy(() => import('./pages/SiteAnnouncements.js'));
|
||||
@@ -348,7 +347,8 @@ export const sidebarGroups = [
|
||||
{ to: '/routes', label: '路由', icon: <svg className="sidebar-item-icon" fill="none" viewBox="0 0 24 24" stroke="currentColor"><path strokeLinecap="round" strokeLinejoin="round" strokeWidth={1.75} d="M8 7h12m0 0l-4-4m4 4l-4 4m0 6H4m0 0l4 4m-4-4l4-4" /></svg> },
|
||||
{ to: '/logs', label: '使用日志', icon: <svg className="sidebar-item-icon" fill="none" viewBox="0 0 24 24" stroke="currentColor"><path strokeLinecap="round" strokeLinejoin="round" strokeWidth={1.75} d="M9 5H7a2 2 0 00-2 2v12a2 2 0 002 2h10a2 2 0 002-2V7a2 2 0 00-2-2h-2M9 5a2 2 0 002 2h2a2 2 0 002-2M9 5a2 2 0 012-2h2a2 2 0 012 2" /></svg> },
|
||||
{ to: '/picture-agent', label: '生图助手', icon: <svg className="sidebar-item-icon" fill="none" viewBox="0 0 24 24" stroke="currentColor"><path strokeLinecap="round" strokeLinejoin="round" strokeWidth={1.5} d="M4 16l4.586-4.586a2 2 0 012.828 0L16 16m-2-2l1.586-1.586a2 2 0 012.828 0L20 14m-5-6h.01M5 4h14a2 2 0 012 2v12a2 2 0 01-2 2H5a2 2 0 01-2-2V6a2 2 0 012-2z" /></svg> },
|
||||
{ to: '/video-agent', label: '视频工作台', icon: <svg className="sidebar-item-icon" fill="none" viewBox="0 0 24 24" stroke="currentColor"><path strokeLinecap="round" strokeLinejoin="round" strokeWidth={1.5} d="M15 10l4.553-2.276A1 1 0 0121 8.618v6.764a1 1 0 01-1.447.894L15 14M5 18h8a2 2 0 002-2V8a2 2 0 00-2-2H5a2 2 0 00-2 2v8a2 2 0 002 2z" /></svg> },
|
||||
{ to: '/digital-human', label: '数字人', icon: <svg className="sidebar-item-icon" fill="none" viewBox="0 0 24 24" stroke="currentColor"><path strokeLinecap="round" strokeLinejoin="round" strokeWidth={1.5} d="M12 12a4 4 0 100-8 4 4 0 000 8z" /><path strokeLinecap="round" strokeLinejoin="round" strokeWidth={1.5} d="M4 21a8 8 0 0116 0M17 9l3 2v4l-3 2M7 9l-3 2v4l3 2" /></svg> },
|
||||
{ to: '/workflows', label: '工作流', icon: <svg className="sidebar-item-icon" fill="none" viewBox="0 0 24 24" stroke="currentColor"><path strokeLinecap="round" strokeLinejoin="round" strokeWidth={1.5} d="M4 6h5v5H4zM15 4h5v5h-5zM15 15h5v5h-5zM9 8h6M17.5 9v6M9 11l6 6" /></svg> },
|
||||
{ to: '/comfyui-agent', label: '视频助手', icon: <svg className="sidebar-item-icon" fill="none" viewBox="0 0 24 24" stroke="currentColor"><path strokeLinecap="round" strokeLinejoin="round" strokeWidth={1.5} d="M14.752 11.168l-3.197-2.132A1 1 0 0010 9.87v4.263a1 1 0 001.555.832l3.197-2.132a1 1 0 000-1.664z" /><path strokeLinecap="round" strokeLinejoin="round" strokeWidth={1.5} d="M21 12a9 9 0 11-18 0 9 9 0 0118 0z" /></svg> },
|
||||
{ to: '/workspace', label: '工作空间', icon: <svg className="sidebar-item-icon" fill="none" viewBox="0 0 24 24" stroke="currentColor"><path strokeLinecap="round" strokeLinejoin="round" strokeWidth={1.5} d="M4 7a2 2 0 012-2h4l2 2h6a2 2 0 012 2v8a2 2 0 01-2 2H6a2 2 0 01-2-2V7z" /><path strokeLinecap="round" strokeLinejoin="round" strokeWidth={1.5} d="M8 13h8M8 16h5" /></svg> },
|
||||
],
|
||||
@@ -369,6 +369,7 @@ export const userSidebarGroups = [
|
||||
label: '概览',
|
||||
items: [
|
||||
{ to: '/user/dashboard', label: '个人中心', icon: <svg className="sidebar-item-icon" fill="none" viewBox="0 0 24 24" stroke="currentColor"><path strokeLinecap="round" strokeLinejoin="round" strokeWidth={1.75} d="M16 7a4 4 0 11-8 0 4 4 0 018 0zM12 14a7 7 0 00-7 7h14a7 7 0 00-7-7z" /></svg> },
|
||||
{ to: '/user/models', label: '模型广场', icon: <svg className="sidebar-item-icon" fill="none" viewBox="0 0 24 24" stroke="currentColor"><path strokeLinecap="round" strokeLinejoin="round" strokeWidth={1.75} d="M4 6h16M4 10h10M4 14h16M4 18h8" /></svg> },
|
||||
{ to: '/user/usage', label: '使用记录', icon: <svg className="sidebar-item-icon" fill="none" viewBox="0 0 24 24" stroke="currentColor"><path strokeLinecap="round" strokeLinejoin="round" strokeWidth={1.75} d="M9 5H7a2 2 0 00-2 2v12a2 2 0 002 2h10a2 2 0 002-2V7a2 2 0 00-2-2h-2M9 5a2 2 0 002 2h2a2 2 0 002-2M9 5a2 2 0 012-2h2a2 2 0 012 2" /></svg> },
|
||||
{ to: '/user/announcements', label: '用户公告', icon: <svg className="sidebar-item-icon" fill="none" viewBox="0 0 24 24" stroke="currentColor"><path strokeLinecap="round" strokeLinejoin="round" strokeWidth={1.75} d="M7 8h10M7 12h10M7 16h6M5 4h14a2 2 0 012 2v12a2 2 0 01-2 2H5a2 2 0 01-2-2V6a2 2 0 012-2z" /></svg> },
|
||||
{ to: '/user/about', label: '关于', icon: <svg className="sidebar-item-icon" fill="none" viewBox="0 0 24 24" stroke="currentColor"><path strokeLinecap="round" strokeLinejoin="round" strokeWidth={1.75} d="M13 16h-1v-4h-1m1-4h.01M12 3a9 9 0 110 18 9 9 0 010-18z" /></svg> },
|
||||
@@ -378,7 +379,8 @@ export const userSidebarGroups = [
|
||||
label: '创作应用',
|
||||
items: [
|
||||
{ to: '/user/picture-agent', label: '生图助手', icon: <svg className="sidebar-item-icon" fill="none" viewBox="0 0 24 24" stroke="currentColor"><path strokeLinecap="round" strokeLinejoin="round" strokeWidth={1.5} d="M4 16l4.586-4.586a2 2 0 012.828 0L16 16m-2-2l1.586-1.586a2 2 0 012.828 0L20 14m-5-6h.01M5 4h14a2 2 0 012 2v12a2 2 0 01-2 2H5a2 2 0 01-2-2V6a2 2 0 012-2z" /></svg> },
|
||||
{ to: '/user/video-agent', label: '视频工作台', icon: <svg className="sidebar-item-icon" fill="none" viewBox="0 0 24 24" stroke="currentColor"><path strokeLinecap="round" strokeLinejoin="round" strokeWidth={1.5} d="M15 10l4.553-2.276A1 1 0 0121 8.618v6.764a1 1 0 01-1.447.894L15 14M5 18h8a2 2 0 002-2V8a2 2 0 00-2-2H5a2 2 0 00-2 2v8a2 2 0 002 2z" /></svg> },
|
||||
{ to: '/user/digital-human', label: '数字人', icon: <svg className="sidebar-item-icon" fill="none" viewBox="0 0 24 24" stroke="currentColor"><path strokeLinecap="round" strokeLinejoin="round" strokeWidth={1.5} d="M12 12a4 4 0 100-8 4 4 0 000 8z" /><path strokeLinecap="round" strokeLinejoin="round" strokeWidth={1.5} d="M4 21a8 8 0 0116 0M17 9l3 2v4l-3 2M7 9l-3 2v4l3 2" /></svg> },
|
||||
{ to: '/user/workflows', label: '工作流', icon: <svg className="sidebar-item-icon" fill="none" viewBox="0 0 24 24" stroke="currentColor"><path strokeLinecap="round" strokeLinejoin="round" strokeWidth={1.5} d="M4 6h5v5H4zM15 4h5v5h-5zM15 15h5v5h-5zM9 8h6M17.5 9v6M9 11l6 6" /></svg> },
|
||||
{ to: '/user/comfyui-agent', label: '视频助手', icon: <svg className="sidebar-item-icon" fill="none" viewBox="0 0 24 24" stroke="currentColor"><path strokeLinecap="round" strokeLinejoin="round" strokeWidth={1.5} d="M14.752 11.168l-3.197-2.132A1 1 0 0010 9.87v4.263a1 1 0 001.555.832l3.197-2.132a1 1 0 000-1.664z" /><path strokeLinecap="round" strokeLinejoin="round" strokeWidth={1.5} d="M21 12a9 9 0 11-18 0 9 9 0 0118 0z" /></svg> },
|
||||
{ to: '/user/workspace', label: '工作空间', icon: <svg className="sidebar-item-icon" fill="none" viewBox="0 0 24 24" stroke="currentColor"><path strokeLinecap="round" strokeLinejoin="round" strokeWidth={1.5} d="M4 7a2 2 0 012-2h4l2 2h6a2 2 0 012 2v8a2 2 0 01-2 2H6a2 2 0 01-2-2V7z" /><path strokeLinecap="round" strokeLinejoin="round" strokeWidth={1.5} d="M8 13h8M8 16h5" /></svg> },
|
||||
],
|
||||
@@ -388,7 +390,7 @@ const topNavItems = [
|
||||
{ label: '控制台', to: '/' },
|
||||
{ label: '模型广场', to: '/models' },
|
||||
{ label: '生图助手', to: '/picture-agent' },
|
||||
{ label: '视频工作台', to: '/video-agent' },
|
||||
{ label: '工作流', to: '/workflows' },
|
||||
{ label: '视频助手', to: '/comfyui-agent' },
|
||||
{ label: '工作空间', to: '/workspace' },
|
||||
{ label: '关于', to: '/about' },
|
||||
@@ -535,12 +537,6 @@ function AppShell() {
|
||||
const timer = setInterval(check, 60_000);
|
||||
return () => clearInterval(timer);
|
||||
}, [authed, toast]);
|
||||
useEffect(() => {
|
||||
if (!authed) return;
|
||||
if (localStorage.getItem(FIRST_USE_DOC_REMINDER_KEY)) return;
|
||||
localStorage.setItem(FIRST_USE_DOC_REMINDER_KEY, '1');
|
||||
toast.info(`${t('首次使用建议先阅读站点文档：')}${SITE_DOCS_URL}`);
|
||||
}, [authed, t, toast]);
|
||||
useEffect(() => {
|
||||
if (!userAuthed || !location.pathname.startsWith('/user/') || location.pathname === '/user/login') {
|
||||
setAnnouncementPopupRows([]);
|
||||
@@ -677,8 +673,9 @@ function AppShell() {
|
||||
const currentSidebarGroups = isUserPage ? userSidebarGroups : sidebarGroups;
|
||||
const currentTopNavItems = isUserPage ? [
|
||||
{ label: '个人中心', to: '/user/dashboard' },
|
||||
{ label: '模型广场', to: '/user/models' },
|
||||
{ label: '用户公告', to: '/user/announcements' },
|
||||
{ label: '视频工作台', to: '/user/video-agent' },
|
||||
{ label: '工作流', to: '/user/workflows' },
|
||||
{ label: '视频助手', to: '/user/comfyui-agent' },
|
||||
{ label: '工作空间', to: '/user/workspace' },
|
||||
{ label: '关于', to: '/user/about' },
|
||||
@@ -917,7 +914,9 @@ function AppShell() {
|
||||
<Route path="/settings/notify" element={<NotificationSettings />} />
|
||||
<Route path="/models" element={<Models />} />
|
||||
<Route path="/picture-agent" element={<PictureAgent />} />
|
||||
<Route path="/video-agent" element={<ComfyUI />} />
|
||||
<Route path="/digital-human" element={<DigitalHuman />} />
|
||||
<Route path="/workflows" element={<Workflows />} />
|
||||
<Route path="/video-agent" element={<Navigate to="/workflows" replace />} />
|
||||
<Route path="/comfyui-agent" element={<ComfyUIAgent />} />
|
||||
<Route path="/workspace" element={<Workspace />} />
|
||||
<Route path="/about" element={<About />} />
|
||||
@@ -928,12 +927,15 @@ function AppShell() {
|
||||
path="/user/dashboard"
|
||||
element={<UserDashboard onLogout={() => setUserAuthed(false)} />}
|
||||
/>
|
||||
<Route path="/user/models" element={<Models />} />
|
||||
<Route path="/user/dashboard/keys" element={<UserKeys />} />
|
||||
<Route path="/user/dashboard/profile" element={<UserProfile />} />
|
||||
<Route path="/user/usage" element={<UsageRecords />} />
|
||||
<Route path="/user/announcements" element={<UserAnnouncements />} />
|
||||
<Route path="/user/picture-agent" element={<PictureAgent />} />
|
||||
<Route path="/user/video-agent" element={<ComfyUI />} />
|
||||
<Route path="/user/digital-human" element={<DigitalHuman />} />
|
||||
<Route path="/user/workflows" element={<Workflows />} />
|
||||
<Route path="/user/video-agent" element={<Navigate to="/user/workflows" replace />} />
|
||||
<Route path="/user/comfyui-agent" element={<ComfyUIAgent />} />
|
||||
<Route path="/user/workspace" element={<Workspace />} />
|
||||
<Route path="/user/about" element={<About />} />
|
||||
|
||||
+115
-2
@@ -55,6 +55,17 @@ async function extractResponseErrorMessage(res: Response): Promise<string> {
|
||||
return message;
|
||||
}
|
||||
|
||||
async function parseJsonResponse<T = any>(res: Response): Promise<T> {
|
||||
if (res.status === 204 || res.status === 205) {
|
||||
return { success: true } as T;
|
||||
}
|
||||
const text = await res.text();
|
||||
if (!text.trim()) {
|
||||
return { success: true } as T;
|
||||
}
|
||||
return JSON.parse(text) as T;
|
||||
}
|
||||
|
||||
function parseContentDispositionFilename(
|
||||
headerValue: string | null,
|
||||
): string | null {
|
||||
@@ -245,7 +256,7 @@ async function userRequest<T = any>(
|
||||
if (!res.ok) {
|
||||
throw new Error(await extractResponseErrorMessage(res));
|
||||
}
|
||||
return res.json() as Promise<T>;
|
||||
return parseJsonResponse<T>(res);
|
||||
}
|
||||
|
||||
async function request<T = any>(
|
||||
@@ -256,7 +267,7 @@ async function request<T = any>(
|
||||
if (!res.ok) {
|
||||
throw new Error(await extractResponseErrorMessage(res));
|
||||
}
|
||||
return res.json() as Promise<T>;
|
||||
return parseJsonResponse<T>(res);
|
||||
}
|
||||
|
||||
function unwrapArrayResponse(value: any, keys: string[] = ["items", "data"]): any[] {
|
||||
@@ -1081,6 +1092,56 @@ export type WorkspaceAssetQuery = {
|
||||
pageSize?: number;
|
||||
};
|
||||
|
||||
export type WorkflowSummary = {
|
||||
id: number;
|
||||
publicId: string;
|
||||
userId?: number | null;
|
||||
name: string;
|
||||
description?: string | null;
|
||||
graph: any;
|
||||
isTemplate?: boolean;
|
||||
visibility: string;
|
||||
createdAt: string;
|
||||
updatedAt: string;
|
||||
};
|
||||
|
||||
export type WorkflowNodeRun = {
|
||||
id: number;
|
||||
nodeId: string;
|
||||
nodeType: string;
|
||||
status: string;
|
||||
inputs: any;
|
||||
outputs?: any;
|
||||
errorMessage?: string | null;
|
||||
startedAt?: string | null;
|
||||
completedAt?: string | null;
|
||||
createdAt: string;
|
||||
};
|
||||
|
||||
export type WorkflowRunSummary = {
|
||||
id: number;
|
||||
publicId: string;
|
||||
workflowId: number;
|
||||
workflowPublicId?: string | null;
|
||||
userId?: number | null;
|
||||
status: string;
|
||||
inputOverrides: any;
|
||||
finalOutputs?: any;
|
||||
errorMessage?: string | null;
|
||||
startedAt?: string | null;
|
||||
completedAt?: string | null;
|
||||
createdAt: string;
|
||||
nodeRuns: WorkflowNodeRun[];
|
||||
};
|
||||
|
||||
export type WorkflowSavePayload = {
|
||||
name?: string;
|
||||
description?: string | null;
|
||||
graph?: any;
|
||||
visibility?: string;
|
||||
isTemplate?: boolean;
|
||||
};
|
||||
|
||||
export const api = {
|
||||
// Sites
|
||||
getSites: async () => {
|
||||
@@ -1939,6 +2000,26 @@ export const api = {
|
||||
request<WorkspaceAssetsResponse>(`/api/workspace/assets${buildQueryString(params)}`),
|
||||
deleteWorkspaceAsset: (publicId: string) =>
|
||||
request(`/api/workspace/assets/${encodeURIComponent(publicId)}`, { method: "DELETE" }),
|
||||
getWorkflowNodeTypes: () =>
|
||||
request<{ success: boolean; nodeTypes: any[] }>("/api/workflow/node-types"),
|
||||
getWorkflows: (params?: { include_public?: boolean }) =>
|
||||
request<{ success: boolean; items: WorkflowSummary[] }>(`/api/workflows${buildQueryString(params)}`),
|
||||
createWorkflow: (data: WorkflowSavePayload) =>
|
||||
request<{ success: boolean; workflow: WorkflowSummary }>("/api/workflows", { method: "POST", body: JSON.stringify(data) }),
|
||||
updateWorkflow: (publicId: string, data: WorkflowSavePayload) =>
|
||||
request<{ success: boolean; workflow: WorkflowSummary }>(`/api/workflows/${encodeURIComponent(publicId)}`, { method: "PUT", body: JSON.stringify(data) }),
|
||||
deleteWorkflow: (publicId: string) =>
|
||||
request(`/api/workflows/${encodeURIComponent(publicId)}`, { method: "DELETE" }),
|
||||
runWorkflow: (publicId: string, inputs?: any) =>
|
||||
request<{ success: boolean; runId: string; run: WorkflowRunSummary }>(`/api/workflows/${encodeURIComponent(publicId)}/run`, {
|
||||
method: "POST",
|
||||
body: JSON.stringify({ inputs: inputs || {} }),
|
||||
timeoutMs: 180_000,
|
||||
}),
|
||||
getWorkflowRuns: (params?: { limit?: number }) =>
|
||||
request<{ success: boolean; items: WorkflowRunSummary[] }>(`/api/workflow-runs${buildQueryString(params)}`),
|
||||
getWorkflowRun: (publicId: string) =>
|
||||
request<{ success: boolean; run: WorkflowRunSummary }>(`/api/workflow-runs/${encodeURIComponent(publicId)}`),
|
||||
};
|
||||
export const userApi = {
|
||||
register: (data: { username: string; email: string; password: string }) =>
|
||||
@@ -1960,6 +2041,18 @@ export const userApi = {
|
||||
getLogs: (params?: { limit?: number; offset?: number; page?: number; page_size?: number; q?: string; status?: string; date_from?: string; date_to?: string }) =>
|
||||
userRequest(`/api/user/logs${buildQueryString(params)}`),
|
||||
getSummary: () => userRequest("/api/user/summary"),
|
||||
getModelsMarketplace: (options?: {
|
||||
refresh?: boolean;
|
||||
includePricing?: boolean;
|
||||
}) => {
|
||||
const params = new URLSearchParams();
|
||||
if (options?.refresh) params.set("refresh", "1");
|
||||
if (options?.includePricing) params.set("includePricing", "1");
|
||||
const query = params.toString();
|
||||
return userRequest(`/api/user/models/marketplace${query ? `?${query}` : ""}`, {
|
||||
timeoutMs: options?.refresh ? 45_000 : 15_000,
|
||||
});
|
||||
},
|
||||
getAnnouncements: (params?: string) =>
|
||||
userRequest(`/api/user/announcements${params ? "?" + params : ""}`),
|
||||
markAnnouncementRead: (id: number) =>
|
||||
@@ -1970,4 +2063,24 @@ export const userApi = {
|
||||
userRequest<WorkspaceAssetsResponse>(`/api/user/workspace/assets${buildQueryString(params)}`),
|
||||
deleteWorkspaceAsset: (publicId: string) =>
|
||||
userRequest(`/api/user/workspace/assets/${encodeURIComponent(publicId)}`, { method: "DELETE" }),
|
||||
getWorkflowNodeTypes: () =>
|
||||
userRequest<{ success: boolean; nodeTypes: any[] }>("/api/user/workflow/node-types"),
|
||||
getWorkflows: (params?: { include_public?: boolean }) =>
|
||||
userRequest<{ success: boolean; items: WorkflowSummary[] }>(`/api/user/workflows${buildQueryString(params)}`),
|
||||
createWorkflow: (data: WorkflowSavePayload) =>
|
||||
userRequest<{ success: boolean; workflow: WorkflowSummary }>("/api/user/workflows", { method: "POST", body: JSON.stringify(data) }),
|
||||
updateWorkflow: (publicId: string, data: WorkflowSavePayload) =>
|
||||
userRequest<{ success: boolean; workflow: WorkflowSummary }>(`/api/user/workflows/${encodeURIComponent(publicId)}`, { method: "PUT", body: JSON.stringify(data) }),
|
||||
deleteWorkflow: (publicId: string) =>
|
||||
userRequest(`/api/user/workflows/${encodeURIComponent(publicId)}`, { method: "DELETE" }),
|
||||
runWorkflow: (publicId: string, inputs?: any) =>
|
||||
userRequest<{ success: boolean; runId: string; run: WorkflowRunSummary }>(`/api/user/workflows/${encodeURIComponent(publicId)}/run`, {
|
||||
method: "POST",
|
||||
body: JSON.stringify({ inputs: inputs || {} }),
|
||||
timeoutMs: 180_000,
|
||||
}),
|
||||
getWorkflowRuns: (params?: { limit?: number }) =>
|
||||
userRequest<{ success: boolean; items: WorkflowRunSummary[] }>(`/api/user/workflow-runs${buildQueryString(params)}`),
|
||||
getWorkflowRun: (publicId: string) =>
|
||||
userRequest<{ success: boolean; run: WorkflowRunSummary }>(`/api/user/workflow-runs/${encodeURIComponent(publicId)}`),
|
||||
};
|
||||
|
||||
@@ -43,6 +43,10 @@ const zhToEn: Record<string, string> = {
|
||||
'仪表盘': 'Dashboard',
|
||||
'站点': 'Sites',
|
||||
'账号': 'Accounts',
|
||||
'连接管理': 'Connection Management',
|
||||
'用户公告': 'User Announcements',
|
||||
'OAuth 管理': 'OAuth Management',
|
||||
'下游密钥': 'Downstream Keys',
|
||||
'未关联站点': 'Unlinked Site',
|
||||
'余额': 'Balance',
|
||||
'令牌管理': 'Account Tokens',
|
||||
@@ -57,11 +61,20 @@ const zhToEn: Record<string, string> = {
|
||||
'使用日志': 'Usage Logs',
|
||||
'暂无使用日志': 'No Usage Logs',
|
||||
'可用性监控': 'Availability Monitor',
|
||||
'生图助手': 'Picture Agent',
|
||||
'数字人': 'Digital Human',
|
||||
'工作流': 'Workflow',
|
||||
'视频助手': 'Video Assistant',
|
||||
'工作空间': 'Workspace',
|
||||
'系统': 'System',
|
||||
'设置': 'Settings',
|
||||
'程序日志': 'System Logs',
|
||||
'导入/导出': 'Import/Export',
|
||||
'通知设置': 'Notification Settings',
|
||||
'概览': 'Overview',
|
||||
'个人中心': 'Profile',
|
||||
'使用记录': 'Usage Records',
|
||||
'创作应用': 'Creative Apps',
|
||||
'暂无通知': 'No Notifications',
|
||||
'模型广场': 'Model Marketplace',
|
||||
'模型广场刷新进行中': 'Marketplace refresh in progress',
|
||||
|
||||
+524
-66
@@ -2887,6 +2887,8 @@ textarea {
|
||||
display: flex;
|
||||
flex-direction: column;
|
||||
gap: 14px;
|
||||
min-width: 0;
|
||||
overflow: hidden;
|
||||
}
|
||||
|
||||
.picture-agent-panel-head {
|
||||
@@ -3115,6 +3117,354 @@ textarea {
|
||||
margin-top: auto;
|
||||
}
|
||||
|
||||
.picture-agent-job-summary {
|
||||
display: flex;
|
||||
align-items: center;
|
||||
justify-content: space-between;
|
||||
gap: 10px;
|
||||
padding: 10px 12px;
|
||||
border: 1px solid var(--color-border-light);
|
||||
border-radius: var(--radius-sm);
|
||||
background: color-mix(in srgb, var(--color-bg) 56%, var(--color-bg-card));
|
||||
min-width: 0;
|
||||
}
|
||||
|
||||
.picture-agent-job-summary div {
|
||||
display: grid;
|
||||
gap: 2px;
|
||||
min-width: 0;
|
||||
}
|
||||
|
||||
.picture-agent-job-summary strong {
|
||||
font-size: 13px;
|
||||
}
|
||||
|
||||
.picture-agent-job-summary span {
|
||||
color: var(--color-text-muted);
|
||||
font-size: 12px;
|
||||
line-height: 1.35;
|
||||
overflow: hidden;
|
||||
text-overflow: ellipsis;
|
||||
white-space: nowrap;
|
||||
}
|
||||
|
||||
.picture-agent-job-summary .btn {
|
||||
flex-shrink: 0;
|
||||
min-height: 30px;
|
||||
padding: 0 10px;
|
||||
font-size: 12px;
|
||||
}
|
||||
|
||||
.picture-agent-job-panel {
|
||||
display: grid;
|
||||
gap: 10px;
|
||||
padding: 12px;
|
||||
border: 1px solid var(--color-border-light);
|
||||
border-radius: var(--radius-sm);
|
||||
background: color-mix(in srgb, var(--color-bg) 56%, var(--color-bg-card));
|
||||
min-width: 0;
|
||||
max-width: 100%;
|
||||
overflow: hidden;
|
||||
}
|
||||
|
||||
.picture-agent-job-head,
|
||||
.picture-agent-job-row,
|
||||
.picture-agent-job-meta {
|
||||
display: flex;
|
||||
align-items: center;
|
||||
justify-content: space-between;
|
||||
gap: 8px;
|
||||
min-width: 0;
|
||||
}
|
||||
|
||||
.picture-agent-job-head strong {
|
||||
font-size: 13px;
|
||||
}
|
||||
|
||||
.picture-agent-job-head-copy {
|
||||
display: grid;
|
||||
gap: 2px;
|
||||
min-width: 0;
|
||||
}
|
||||
|
||||
.picture-agent-job-head-copy span {
|
||||
color: var(--color-text-muted);
|
||||
font-size: 11px;
|
||||
line-height: 1.4;
|
||||
}
|
||||
|
||||
.picture-agent-job-head .btn {
|
||||
min-height: 28px;
|
||||
padding: 0 8px;
|
||||
font-size: 12px;
|
||||
}
|
||||
|
||||
.picture-agent-job-actions {
|
||||
display: flex;
|
||||
align-items: center;
|
||||
gap: 6px;
|
||||
flex-wrap: wrap;
|
||||
justify-content: flex-end;
|
||||
}
|
||||
|
||||
.picture-agent-job-empty {
|
||||
color: var(--color-text-muted);
|
||||
font-size: 12px;
|
||||
line-height: 1.55;
|
||||
}
|
||||
|
||||
.picture-agent-job-list {
|
||||
display: grid;
|
||||
gap: 6px;
|
||||
min-width: 0;
|
||||
}
|
||||
|
||||
.picture-agent-job {
|
||||
display: grid;
|
||||
gap: 5px;
|
||||
padding: 8px 9px;
|
||||
border: 1px solid var(--color-border-light);
|
||||
border-radius: 8px;
|
||||
background: var(--color-bg-card);
|
||||
min-width: 0;
|
||||
max-width: 100%;
|
||||
overflow: hidden;
|
||||
}
|
||||
|
||||
.picture-agent-job-row-top {
|
||||
align-items: flex-start;
|
||||
}
|
||||
|
||||
.picture-agent-job-title-wrap {
|
||||
display: grid;
|
||||
gap: 2px;
|
||||
min-width: 0;
|
||||
}
|
||||
|
||||
.picture-agent-job-row strong {
|
||||
min-width: 0;
|
||||
overflow: hidden;
|
||||
display: block;
|
||||
font-size: 12px;
|
||||
line-height: 1.4;
|
||||
overflow-wrap: anywhere;
|
||||
}
|
||||
|
||||
.picture-agent-job-title-wrap span {
|
||||
color: var(--color-text-muted);
|
||||
font-size: 11px;
|
||||
line-height: 1.35;
|
||||
overflow-wrap: anywhere;
|
||||
}
|
||||
|
||||
.picture-agent-job-row span {
|
||||
flex-shrink: 0;
|
||||
color: var(--color-text-muted);
|
||||
font-size: 12px;
|
||||
}
|
||||
|
||||
.picture-agent-job-status {
|
||||
flex-shrink: 0;
|
||||
font-size: 11px;
|
||||
line-height: 1.4;
|
||||
padding: 2px 8px;
|
||||
border-radius: 999px;
|
||||
border: 1px solid var(--color-border-light);
|
||||
background: color-mix(in srgb, var(--color-bg) 72%, transparent);
|
||||
}
|
||||
|
||||
.picture-agent-job-status-running,
|
||||
.picture-agent-job-status-pending {
|
||||
color: var(--color-primary);
|
||||
}
|
||||
|
||||
.picture-agent-job-status-succeeded {
|
||||
color: var(--color-success);
|
||||
}
|
||||
|
||||
.picture-agent-job-status-failed {
|
||||
color: var(--color-danger);
|
||||
}
|
||||
|
||||
.picture-agent-job-status-cancelled {
|
||||
color: var(--color-text-muted);
|
||||
}
|
||||
|
||||
.picture-agent-job-progress-line {
|
||||
display: grid;
|
||||
gap: 4px;
|
||||
}
|
||||
|
||||
.picture-agent-job-progress-labels {
|
||||
display: flex;
|
||||
align-items: center;
|
||||
justify-content: space-between;
|
||||
gap: 8px;
|
||||
min-width: 0;
|
||||
color: var(--color-text-muted);
|
||||
font-size: 11px;
|
||||
line-height: 1.35;
|
||||
}
|
||||
|
||||
.picture-agent-job-progress-labels span,
|
||||
.picture-agent-job-progress-labels em {
|
||||
min-width: 0;
|
||||
overflow: hidden;
|
||||
text-overflow: ellipsis;
|
||||
white-space: nowrap;
|
||||
}
|
||||
|
||||
.picture-agent-job-progress-labels em {
|
||||
font-style: normal;
|
||||
}
|
||||
|
||||
.picture-agent-job-progress {
|
||||
height: 6px;
|
||||
width: 100%;
|
||||
min-width: 0;
|
||||
overflow: hidden;
|
||||
border-radius: 999px;
|
||||
background: color-mix(in srgb, var(--color-border) 52%, transparent);
|
||||
}
|
||||
|
||||
.picture-agent-job-progress span {
|
||||
display: block;
|
||||
height: 100%;
|
||||
border-radius: inherit;
|
||||
background: linear-gradient(90deg, var(--color-primary), var(--color-info));
|
||||
transition: width 0.22s ease;
|
||||
}
|
||||
|
||||
.picture-agent-job-succeeded .picture-agent-job-progress span {
|
||||
background: var(--color-success);
|
||||
}
|
||||
|
||||
.picture-agent-job-failed .picture-agent-job-progress span {
|
||||
background: var(--color-danger);
|
||||
}
|
||||
|
||||
.picture-agent-job-cancelled .picture-agent-job-progress span {
|
||||
background: var(--color-text-muted);
|
||||
}
|
||||
|
||||
.picture-agent-job-meta {
|
||||
align-items: flex-start;
|
||||
color: var(--color-text-muted);
|
||||
font-size: 11px;
|
||||
line-height: 1.45;
|
||||
}
|
||||
|
||||
.picture-agent-job-meta span,
|
||||
.picture-agent-job-meta em {
|
||||
min-width: 0;
|
||||
overflow: hidden;
|
||||
text-overflow: ellipsis;
|
||||
white-space: nowrap;
|
||||
}
|
||||
|
||||
.picture-agent-job-meta em {
|
||||
color: var(--color-danger);
|
||||
font-style: normal;
|
||||
}
|
||||
|
||||
.picture-agent-job-hint {
|
||||
color: var(--color-warning);
|
||||
font-size: 11px;
|
||||
line-height: 1.45;
|
||||
overflow-wrap: anywhere;
|
||||
}
|
||||
|
||||
.picture-agent-job-terminal {
|
||||
color: var(--color-text-muted);
|
||||
font-size: 11px;
|
||||
line-height: 1.45;
|
||||
}
|
||||
|
||||
.picture-agent-job-cancel {
|
||||
justify-self: end;
|
||||
border: 1px solid var(--color-border-light);
|
||||
border-radius: 999px;
|
||||
padding: 3px 9px;
|
||||
background: transparent;
|
||||
color: var(--color-text-muted);
|
||||
font-size: 11px;
|
||||
cursor: pointer;
|
||||
}
|
||||
|
||||
.picture-agent-job-cancel:hover:not(:disabled) {
|
||||
color: var(--color-danger);
|
||||
border-color: color-mix(in srgb, var(--color-danger) 46%, var(--color-border-light));
|
||||
}
|
||||
|
||||
.picture-agent-job-cancel:disabled {
|
||||
cursor: not-allowed;
|
||||
opacity: 0.58;
|
||||
}
|
||||
|
||||
.picture-agent-job-dialog-backdrop {
|
||||
position: fixed;
|
||||
inset: 0;
|
||||
z-index: 80;
|
||||
display: grid;
|
||||
place-items: center;
|
||||
padding: 24px;
|
||||
background: rgba(15, 23, 42, 0.42);
|
||||
backdrop-filter: blur(2px);
|
||||
}
|
||||
|
||||
.picture-agent-job-dialog {
|
||||
width: min(820px, calc(100vw - 40px));
|
||||
max-height: min(760px, calc(100vh - 56px));
|
||||
display: grid;
|
||||
grid-template-rows: auto minmax(0, 1fr);
|
||||
gap: 12px;
|
||||
padding: 16px;
|
||||
border: 1px solid var(--color-border);
|
||||
border-radius: var(--radius-sm);
|
||||
background: var(--color-bg-card);
|
||||
box-shadow: var(--shadow-lg);
|
||||
overflow: hidden;
|
||||
}
|
||||
|
||||
.picture-agent-job-dialog-head {
|
||||
display: flex;
|
||||
align-items: flex-start;
|
||||
justify-content: space-between;
|
||||
gap: 12px;
|
||||
padding-bottom: 12px;
|
||||
border-bottom: 1px solid var(--color-border-light);
|
||||
}
|
||||
|
||||
.picture-agent-job-dialog-head > div:first-child {
|
||||
display: grid;
|
||||
gap: 4px;
|
||||
min-width: 0;
|
||||
}
|
||||
|
||||
.picture-agent-job-dialog-head strong {
|
||||
font-size: 16px;
|
||||
}
|
||||
|
||||
.picture-agent-job-dialog-head span {
|
||||
color: var(--color-text-muted);
|
||||
font-size: 12px;
|
||||
}
|
||||
|
||||
.picture-agent-job-list-dialog {
|
||||
max-height: 100%;
|
||||
overflow-y: auto;
|
||||
padding-right: 4px;
|
||||
}
|
||||
|
||||
.picture-agent-job-dialog .picture-agent-job {
|
||||
padding: 12px;
|
||||
}
|
||||
|
||||
.picture-agent-job-dialog .picture-agent-job-row strong {
|
||||
max-height: 5.6em;
|
||||
overflow-y: auto;
|
||||
}
|
||||
|
||||
.picture-agent-canvas {
|
||||
display: grid;
|
||||
grid-template-rows: auto minmax(0, 1fr) auto;
|
||||
@@ -3146,9 +3496,10 @@ textarea {
|
||||
|
||||
.picture-agent-template-list {
|
||||
display: flex;
|
||||
gap: 10px;
|
||||
gap: 8px;
|
||||
overflow-x: auto;
|
||||
padding-bottom: 2px;
|
||||
padding: 2px 2px 6px;
|
||||
scrollbar-width: thin;
|
||||
}
|
||||
|
||||
.picture-agent-template-empty {
|
||||
@@ -3166,13 +3517,13 @@ textarea {
|
||||
}
|
||||
|
||||
.picture-agent-template-card {
|
||||
width: 156px;
|
||||
min-height: 168px;
|
||||
width: 142px;
|
||||
min-height: 154px;
|
||||
flex: 0 0 auto;
|
||||
display: grid;
|
||||
grid-template-rows: 112px auto;
|
||||
gap: 8px;
|
||||
padding: 8px;
|
||||
grid-template-rows: 102px auto;
|
||||
gap: 7px;
|
||||
padding: 7px;
|
||||
border: 1px solid var(--color-border);
|
||||
border-radius: var(--radius-sm);
|
||||
background: var(--color-bg-card);
|
||||
@@ -3188,7 +3539,7 @@ textarea {
|
||||
|
||||
.picture-agent-template-card img {
|
||||
width: 100%;
|
||||
height: 112px;
|
||||
height: 102px;
|
||||
display: block;
|
||||
object-fit: cover;
|
||||
border-radius: 6px;
|
||||
@@ -3209,12 +3560,12 @@ textarea {
|
||||
}
|
||||
|
||||
.picture-agent-template-card strong {
|
||||
font-size: 13px;
|
||||
font-size: 12px;
|
||||
}
|
||||
|
||||
.picture-agent-template-card em {
|
||||
color: var(--color-text-muted);
|
||||
font-size: 12px;
|
||||
font-size: 11px;
|
||||
font-style: normal;
|
||||
}
|
||||
|
||||
@@ -4287,11 +4638,9 @@ textarea {
|
||||
display: flex;
|
||||
align-items: center;
|
||||
justify-content: center;
|
||||
padding: 32px;
|
||||
padding: 30px;
|
||||
background:
|
||||
radial-gradient(circle at top left, rgba(59, 130, 246, 0.18), transparent 34%),
|
||||
radial-gradient(circle at right 18%, rgba(34, 211, 238, 0.16), transparent 30%),
|
||||
linear-gradient(180deg, #eef4fb 0%, #e9eef7 100%);
|
||||
linear-gradient(135deg, rgba(219, 234, 254, 0.92) 0%, rgba(248, 250, 252, 0.98) 44%, rgba(224, 242, 254, 0.88) 100%);
|
||||
}
|
||||
|
||||
[data-theme="dark"] .login-shell {
|
||||
@@ -4302,17 +4651,18 @@ textarea {
|
||||
}
|
||||
|
||||
.login-surface {
|
||||
width: min(1240px, 100%);
|
||||
width: min(1280px, 100%);
|
||||
display: grid;
|
||||
grid-template-columns: minmax(0, 1.16fr) minmax(360px, 440px);
|
||||
gap: 24px;
|
||||
padding: 20px;
|
||||
grid-template-columns: minmax(620px, 1fr) minmax(360px, 420px);
|
||||
gap: 18px;
|
||||
padding: 18px;
|
||||
background:
|
||||
linear-gradient(145deg, rgba(255, 255, 255, 0.78), rgba(244, 248, 253, 0.88));
|
||||
border: 1px solid rgba(255, 255, 255, 0.58);
|
||||
border-radius: 36px;
|
||||
linear-gradient(145deg, rgba(255, 255, 255, 0.9), rgba(244, 248, 253, 0.94));
|
||||
border: 1px solid rgba(255, 255, 255, 0.72);
|
||||
border-radius: 30px;
|
||||
box-shadow:
|
||||
0 24px 60px rgba(148, 163, 184, 0.24),
|
||||
0 30px 72px rgba(59, 130, 246, 0.16),
|
||||
0 18px 44px rgba(15, 23, 42, 0.08),
|
||||
inset 0 1px 0 rgba(255, 255, 255, 0.9);
|
||||
}
|
||||
|
||||
@@ -4327,21 +4677,19 @@ textarea {
|
||||
|
||||
.login-brand-panel {
|
||||
position: relative;
|
||||
min-height: 760px;
|
||||
padding: 56px 56px 44px;
|
||||
min-height: 720px;
|
||||
padding: 54px 58px 42px;
|
||||
display: flex;
|
||||
flex-direction: column;
|
||||
gap: 28px;
|
||||
gap: 30px;
|
||||
color: var(--color-text-primary);
|
||||
border-radius: 30px;
|
||||
border-radius: 24px;
|
||||
overflow: hidden;
|
||||
}
|
||||
|
||||
.login-brand-panel-light {
|
||||
background:
|
||||
radial-gradient(circle at top left, rgba(96, 165, 250, 0.22), transparent 26%),
|
||||
radial-gradient(circle at right 20%, rgba(34, 211, 238, 0.16), transparent 22%),
|
||||
linear-gradient(160deg, #f8fbff 0%, #edf4ff 52%, #e9f0fa 100%);
|
||||
linear-gradient(160deg, rgba(255, 255, 255, 0.96) 0%, rgba(239, 246, 255, 0.94) 46%, rgba(226, 239, 255, 0.9) 100%);
|
||||
}
|
||||
|
||||
.login-brand-panel::before {
|
||||
@@ -4369,7 +4717,8 @@ textarea {
|
||||
.login-brand-header {
|
||||
display: flex;
|
||||
align-items: center;
|
||||
gap: 20px;
|
||||
gap: 22px;
|
||||
z-index: 1;
|
||||
}
|
||||
|
||||
.login-brand-summary {
|
||||
@@ -4402,9 +4751,9 @@ textarea {
|
||||
}
|
||||
|
||||
.brand-mark-frame-hero {
|
||||
width: 118px;
|
||||
height: 118px;
|
||||
border-radius: 36px;
|
||||
width: 124px;
|
||||
height: 124px;
|
||||
border-radius: 32px;
|
||||
}
|
||||
|
||||
.brand-mark-canvas {
|
||||
@@ -4429,9 +4778,10 @@ textarea {
|
||||
|
||||
.login-brand-name {
|
||||
font-family: 'Avenir Next', 'Segoe UI Variable Display', 'PingFang SC', 'Noto Sans SC', sans-serif;
|
||||
font-size: clamp(34px, 3.2vw, 46px);
|
||||
font-size: clamp(38px, 3.6vw, 54px);
|
||||
font-weight: 800;
|
||||
letter-spacing: -0.05em;
|
||||
line-height: 1;
|
||||
letter-spacing: 0;
|
||||
}
|
||||
|
||||
.login-brand-kicker {
|
||||
@@ -4444,47 +4794,57 @@ textarea {
|
||||
|
||||
.login-brand-copy-block {
|
||||
display: grid;
|
||||
max-width: 580px;
|
||||
max-width: 620px;
|
||||
z-index: 1;
|
||||
}
|
||||
|
||||
.login-brand-copy {
|
||||
max-width: 42ch;
|
||||
max-width: 44ch;
|
||||
margin: 0;
|
||||
font-size: 18px;
|
||||
line-height: 1.92;
|
||||
font-size: 20px;
|
||||
line-height: 1.82;
|
||||
color: var(--color-text-secondary);
|
||||
}
|
||||
|
||||
.login-compat-line {
|
||||
max-width: 42ch;
|
||||
max-width: 52ch;
|
||||
font-size: 13px;
|
||||
line-height: 1.8;
|
||||
color: #475569;
|
||||
z-index: 1;
|
||||
}
|
||||
|
||||
.login-capability-list {
|
||||
display: grid;
|
||||
gap: 0;
|
||||
max-width: 44rem;
|
||||
max-width: 46rem;
|
||||
margin-top: auto;
|
||||
border-top: 1px solid rgba(148, 163, 184, 0.18);
|
||||
z-index: 1;
|
||||
}
|
||||
|
||||
.login-capability-row {
|
||||
display: grid;
|
||||
grid-template-columns: 36px minmax(0, 1fr);
|
||||
gap: 16px;
|
||||
grid-template-columns: 48px minmax(0, 1fr);
|
||||
gap: 18px;
|
||||
align-items: start;
|
||||
padding: 18px 0;
|
||||
padding: 20px 0;
|
||||
border-bottom: 1px solid rgba(148, 163, 184, 0.16);
|
||||
}
|
||||
|
||||
.login-capability-index {
|
||||
padding-top: 3px;
|
||||
font-size: 11px;
|
||||
width: 38px;
|
||||
height: 38px;
|
||||
display: inline-flex;
|
||||
align-items: center;
|
||||
justify-content: center;
|
||||
font-size: 12px;
|
||||
font-weight: 700;
|
||||
letter-spacing: 0.18em;
|
||||
color: #64748b;
|
||||
letter-spacing: 0;
|
||||
color: #2563eb;
|
||||
border: 1px solid rgba(37, 99, 235, 0.14);
|
||||
border-radius: 12px;
|
||||
background: rgba(255, 255, 255, 0.68);
|
||||
}
|
||||
|
||||
.login-capability-content {
|
||||
@@ -4494,9 +4854,9 @@ textarea {
|
||||
|
||||
.login-capability-title {
|
||||
font-family: 'Avenir Next', 'Segoe UI Variable Display', 'PingFang SC', 'Noto Sans SC', sans-serif;
|
||||
font-size: 18px;
|
||||
font-size: 19px;
|
||||
font-weight: 700;
|
||||
letter-spacing: -0.03em;
|
||||
letter-spacing: 0;
|
||||
color: #0f172a;
|
||||
}
|
||||
|
||||
@@ -4512,6 +4872,7 @@ textarea {
|
||||
display: flex;
|
||||
align-items: center;
|
||||
gap: 12px;
|
||||
z-index: 1;
|
||||
}
|
||||
|
||||
.login-icon-link {
|
||||
@@ -4563,21 +4924,21 @@ textarea {
|
||||
display: flex;
|
||||
align-items: center;
|
||||
justify-content: center;
|
||||
padding: 24px;
|
||||
border-radius: 30px;
|
||||
padding: 22px;
|
||||
border-radius: 24px;
|
||||
background:
|
||||
radial-gradient(circle at top, rgba(125, 211, 252, 0.16), transparent 34%),
|
||||
linear-gradient(180deg, rgba(255, 255, 255, 0.44), rgba(241, 245, 249, 0.58));
|
||||
linear-gradient(180deg, rgba(255, 255, 255, 0.56), rgba(241, 245, 249, 0.7));
|
||||
}
|
||||
|
||||
.login-auth-panel {
|
||||
width: min(100%, 392px);
|
||||
padding: 40px 34px 28px;
|
||||
border-radius: 28px;
|
||||
padding: 38px 34px 28px;
|
||||
border-radius: 22px;
|
||||
border: 1px solid rgba(226, 232, 240, 0.92);
|
||||
background: rgba(255, 255, 255, 0.92);
|
||||
box-shadow:
|
||||
0 28px 50px rgba(148, 163, 184, 0.22),
|
||||
0 24px 44px rgba(37, 99, 235, 0.1),
|
||||
0 10px 22px rgba(15, 23, 42, 0.08),
|
||||
inset 0 1px 0 rgba(255, 255, 255, 0.94);
|
||||
backdrop-filter: blur(10px);
|
||||
}
|
||||
@@ -4594,10 +4955,10 @@ textarea {
|
||||
.login-auth-title {
|
||||
margin: 0 0 10px;
|
||||
font-family: 'Avenir Next', 'Segoe UI Variable Display', 'PingFang SC', 'Noto Sans SC', sans-serif;
|
||||
font-size: 42px;
|
||||
font-size: 38px;
|
||||
font-weight: 800;
|
||||
line-height: 0.96;
|
||||
letter-spacing: -0.05em;
|
||||
line-height: 1.08;
|
||||
letter-spacing: 0;
|
||||
}
|
||||
|
||||
.login-auth-copy {
|
||||
@@ -4610,18 +4971,18 @@ textarea {
|
||||
.login-mode-tabs {
|
||||
display: grid;
|
||||
grid-template-columns: repeat(2, minmax(0, 1fr));
|
||||
gap: 4px;
|
||||
gap: 6px;
|
||||
padding: 4px;
|
||||
margin-bottom: 18px;
|
||||
border: 1px solid color-mix(in srgb, var(--color-primary) 12%, var(--color-border));
|
||||
border-radius: 16px;
|
||||
border-radius: 14px;
|
||||
background: color-mix(in srgb, var(--color-bg) 78%, transparent);
|
||||
}
|
||||
|
||||
.login-mode-tab {
|
||||
min-height: 38px;
|
||||
min-height: 40px;
|
||||
border: 0;
|
||||
border-radius: 12px;
|
||||
border-radius: 10px;
|
||||
background: transparent;
|
||||
color: var(--color-text-muted);
|
||||
font-size: 13px;
|
||||
@@ -4648,7 +5009,7 @@ textarea {
|
||||
width: 100%;
|
||||
padding: 14px 16px;
|
||||
border: 1px solid color-mix(in srgb, var(--color-primary) 18%, var(--color-border));
|
||||
border-radius: 16px;
|
||||
border-radius: 12px;
|
||||
font-size: 14px;
|
||||
outline: none;
|
||||
background: rgba(248, 250, 252, 0.92);
|
||||
@@ -4668,7 +5029,7 @@ textarea {
|
||||
padding: 13px 0;
|
||||
font-size: 14px;
|
||||
font-weight: 700;
|
||||
border-radius: 16px;
|
||||
border-radius: 12px;
|
||||
box-shadow: 0 16px 24px rgba(99, 102, 241, 0.2);
|
||||
}
|
||||
|
||||
@@ -4687,6 +5048,18 @@ textarea {
|
||||
color: var(--color-text-muted);
|
||||
}
|
||||
|
||||
.login-register-form .login-auth-label {
|
||||
margin-bottom: 6px;
|
||||
}
|
||||
|
||||
.login-register-form .login-auth-input {
|
||||
margin-bottom: 12px;
|
||||
}
|
||||
|
||||
.login-register-form .login-auth-submit {
|
||||
margin-top: 2px;
|
||||
}
|
||||
|
||||
.login-shell[data-variant="user"] .login-surface {
|
||||
width: min(560px, 100%);
|
||||
grid-template-columns: minmax(0, 1fr);
|
||||
@@ -4751,6 +5124,35 @@ textarea {
|
||||
max-width: 100%;
|
||||
}
|
||||
|
||||
@media (max-height: 900px) and (min-width: 641px) {
|
||||
.login-shell[data-variant="user"] {
|
||||
align-items: flex-start;
|
||||
overflow-y: auto;
|
||||
}
|
||||
|
||||
.login-shell[data-variant="user"] .login-surface {
|
||||
margin: auto 0;
|
||||
}
|
||||
|
||||
.login-shell[data-variant="user"] .login-brand-panel {
|
||||
padding-block: 20px 14px;
|
||||
}
|
||||
|
||||
.login-shell[data-variant="user"] .brand-mark-frame-hero {
|
||||
width: 64px;
|
||||
height: 64px;
|
||||
border-radius: 20px;
|
||||
}
|
||||
|
||||
.login-shell[data-variant="user"] .login-brand-name {
|
||||
font-size: 30px;
|
||||
}
|
||||
|
||||
.login-shell[data-variant="user"] .login-auth-panel {
|
||||
padding: 30px 30px 24px;
|
||||
}
|
||||
}
|
||||
|
||||
[data-theme="dark"] .login-surface,
|
||||
[data-theme="dark"] .brand-mark-frame,
|
||||
[data-theme="dark"] .brand-mark-canvas,
|
||||
@@ -7788,3 +8190,59 @@ textarea {
|
||||
grid-template-columns: 1fr;
|
||||
}
|
||||
}
|
||||
|
||||
.comfyui-agent-page article p {
|
||||
margin: 0 0 0.72em;
|
||||
}
|
||||
|
||||
.comfyui-agent-page article p:last-child {
|
||||
margin-bottom: 0;
|
||||
}
|
||||
|
||||
.comfyui-agent-page article ul,
|
||||
.comfyui-agent-page article ol {
|
||||
margin: 0.4em 0 0.8em;
|
||||
padding-left: 1.35em;
|
||||
}
|
||||
|
||||
.comfyui-agent-page article blockquote {
|
||||
margin: 0.6em 0;
|
||||
padding: 0.2em 0 0.2em 0.9em;
|
||||
border-left: 3px solid color-mix(in srgb, var(--color-primary) 36%, var(--color-border));
|
||||
color: var(--color-text-secondary);
|
||||
}
|
||||
|
||||
.comfyui-agent-page article pre {
|
||||
max-width: 100%;
|
||||
overflow: auto;
|
||||
margin: 0.75em 0;
|
||||
padding: 10px 12px;
|
||||
border-radius: 8px;
|
||||
background: color-mix(in srgb, var(--color-bg) 82%, #111827);
|
||||
border: 1px solid var(--color-border-light);
|
||||
}
|
||||
|
||||
.comfyui-agent-page article code {
|
||||
font-family: var(--font-mono);
|
||||
font-size: 0.92em;
|
||||
}
|
||||
|
||||
.comfyui-agent-page article :not(pre) > code {
|
||||
padding: 0.12em 0.35em;
|
||||
border-radius: 5px;
|
||||
background: color-mix(in srgb, var(--color-primary) 10%, transparent);
|
||||
}
|
||||
|
||||
.comfyui-agent-page article table {
|
||||
display: block;
|
||||
max-width: 100%;
|
||||
overflow: auto;
|
||||
border-collapse: collapse;
|
||||
margin: 0.8em 0;
|
||||
}
|
||||
|
||||
.comfyui-agent-page article th,
|
||||
.comfyui-agent-page article td {
|
||||
border: 1px solid var(--color-border-light);
|
||||
padding: 6px 8px;
|
||||
}
|
||||
|
||||
@@ -7,8 +7,8 @@ const FEATURES = [
|
||||
{ icon: 'RT', title: '模型路由与容灾', desc: '按模型、账号、站点和通道优先级调度，失败后自动冷却并切换到可用通道。' },
|
||||
{ icon: 'AC', title: '站点与账号管理', desc: '集中维护站点、账号、API Key、OAuth、下游 Key 和普通用户可用模型。' },
|
||||
{ icon: 'MD', title: '模型广场', desc: '聚合各站点模型可用性、账号覆盖、绑定令牌与路由状态，便于快速修路由。' },
|
||||
{ icon: 'CU', title: 'ComfyUI 工作台', desc: '内置多租户 ComfyUI 容器池，支持用 BoosAPI 模型驱动文本、图像和视频节点。' },
|
||||
{ icon: 'AG', title: 'ComfyUI Agent', desc: '面向工作流创作的对话助手，支持附件、多会话、用户 API Key 和模型路由选择。' },
|
||||
{ icon: 'WF', title: '工作流画布', desc: '沿用 ComfyUI 节点画布体验，后续接入 BoosAPI 模型路由、图像、视频、音频和工具节点。' },
|
||||
{ icon: 'AG', title: '视频助手', desc: '面向工作流创作的对话助手，支持附件、多会话、用户 API Key 和模型路由选择。' },
|
||||
{ icon: 'PI', title: '生图助手', desc: '提供提示词模板、比例尺寸控制和图像模型调用，管理员与普通用户均可使用。' },
|
||||
{ icon: 'VD', title: '视频任务模板', desc: '支持火山方舟 Step1 创建视频任务、Step2 查询任务结果，并可复用 BoosAPI 模型路由。' },
|
||||
{ icon: 'CK', title: '签到与余额刷新', desc: '按站点/账号执行签到、余额刷新、健康检查和连接状态追踪。' },
|
||||
@@ -17,18 +17,18 @@ const FEATURES = [
|
||||
|
||||
const TECH_STACK = [
|
||||
{ name: 'Rust', desc: '后端核心服务，提供代理、路由、认证、任务和静态资源服务。' },
|
||||
{ name: 'Axum / Tokio', desc: '异步 HTTP 服务、SSE、后台任务和容器生命周期管理。' },
|
||||
{ name: 'Axum / Tokio', desc: '异步 HTTP 服务、SSE、后台任务和代理执行。' },
|
||||
{ name: 'SQLx', desc: '类型化数据库访问与迁移管理。' },
|
||||
{ name: 'SQLite', desc: '默认嵌入式数据库，适合单机部署和低运维成本场景。' },
|
||||
{ name: 'React / Vite', desc: '管理后台、普通用户面板、模型测试器和 Agent 界面。' },
|
||||
{ name: 'ComfyUI / Docker', desc: '视频工作台和多租户工作流容器运行环境。' },
|
||||
{ name: 'External ComfyUI', desc: '可选外部 ComfyUI 代理；BoosAPI 不再内置 Docker 工作台。' },
|
||||
];
|
||||
|
||||
const ENDPOINTS = [
|
||||
{ path: '/v1/chat/completions', desc: 'OpenAI 兼容对话代理，支持非流式与流式。' },
|
||||
{ path: '/v1/messages', desc: 'Claude Messages API 兼容入口，支持 Anthropic 请求、响应和流式事件。' },
|
||||
{ path: '/v1/responses', desc: 'Responses API 兼容代理，可接入支持该协议的上游。' },
|
||||
{ path: '/v1/images/generations', desc: '图像生成代理，供生图助手和 ComfyUI 图像节点使用。' },
|
||||
{ path: '/v1/images/generations', desc: '图像生成代理，供生图助手和工作流图像节点使用。' },
|
||||
{ path: '/v1/videos', desc: '视频生成任务入口，已适配火山方舟任务创建。' },
|
||||
{ path: '/v1/videos/{id}', desc: '视频任务查询入口，已适配火山方舟任务轮询。' },
|
||||
{ path: '/v1/models', desc: '仅暴露当前可路由、未冷却、账号有效的模型列表。' },
|
||||
@@ -56,10 +56,10 @@ export default function About() {
|
||||
</div>
|
||||
</div>
|
||||
<div style={{ fontSize: 14, color: 'var(--color-text-secondary)', lineHeight: 1.8 }}>
|
||||
{tr('BoosAPI 是一个自托管 AI 模型网关与工作台。它把多站点账号、下游 Key、模型路由、ComfyUI、图像生成、视频任务和普通用户面板收束到同一个 Rust 服务中，让前端、Agent 和 ComfyUI 节点都能使用后台已经配置好的模型。')}
|
||||
{tr('BoosAPI 是一个自托管 AI 模型网关与工作台。它把多站点账号、下游 Key、模型路由、图像生成、视频任务、ComfyUI 风格工作流和普通用户面板收束到同一个 Rust 服务中，让前端、Agent 和工作流节点都能使用后台已经配置好的模型。')}
|
||||
</div>
|
||||
<div style={{ display: 'flex', gap: 8, flexWrap: 'wrap', marginTop: 14 }}>
|
||||
{['Rust Backend', 'OpenAI Compatible', 'ComfyUI', 'Picture Agent', 'Volcengine Video'].map((tag) => (
|
||||
{['Rust Backend', 'OpenAI Compatible', 'ComfyUI Canvas', 'Picture Agent', 'Volcengine Video'].map((tag) => (
|
||||
<span
|
||||
key={tag}
|
||||
style={{
|
||||
@@ -171,7 +171,7 @@ export default function About() {
|
||||
<div className="card animate-slide-up stagger-5" style={{ padding: 22 }}>
|
||||
<h3 style={{ fontSize: 15, fontWeight: 600, marginBottom: 10 }}>{tr('数据与隐私')}</h3>
|
||||
<div style={{ fontSize: 13, color: 'var(--color-text-secondary)', lineHeight: 1.8 }}>
|
||||
{tr('BoosAPI 完全自托管，账号、令牌、路由、日志、普通用户和工作台状态均保存在你的运行环境中。代理请求只在你的服务器、ComfyUI 容器与已配置上游之间流转；内置系统 Key 会被保护，后台列表中只展示脱敏值。')}
|
||||
{tr('BoosAPI 完全自托管，账号、令牌、路由、日志、普通用户和工作台状态均保存在你的运行环境中。代理请求只在你的服务器、可选外部 ComfyUI 与已配置上游之间流转；内置系统 Key 会被保护，后台列表中只展示脱敏值。')}
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
@@ -493,7 +493,37 @@ export default function Accounts() {
|
||||
credentialMode,
|
||||
skipModelFetch: tokenForm.skipModelFetch,
|
||||
});
|
||||
const createdAccountIds = result?.batch
|
||||
? Array.isArray(result?.data)
|
||||
? result.data
|
||||
.map((item: any) => Number(item?.id))
|
||||
.filter((id: number) => Number.isFinite(id) && id > 0)
|
||||
: []
|
||||
: Number(result?.id) > 0
|
||||
? [Number(result.id)]
|
||||
: [];
|
||||
if (result?.batch) {
|
||||
if (
|
||||
activeSegment === "apikey" &&
|
||||
createdAccountIds.length > 0
|
||||
) {
|
||||
try {
|
||||
const syncResults = await Promise.allSettled(
|
||||
createdAccountIds.map((accountId: number) =>
|
||||
api.checkModels(accountId),
|
||||
),
|
||||
);
|
||||
const syncedCount = syncResults.filter(
|
||||
(item) => item.status === "fulfilled",
|
||||
).length;
|
||||
if (syncedCount > 0) {
|
||||
await api.rebuildRoutes(false, false);
|
||||
toast.success("批量 API Key 已自动同步模型并重建路由");
|
||||
}
|
||||
} catch (syncErr: any) {
|
||||
toast.error(syncErr?.message || "批量连接已添加，但自动建路由失败");
|
||||
}
|
||||
}
|
||||
closeAddPanel();
|
||||
const createdCount = Number(result.createdCount) || 0;
|
||||
const failedCount = Number(result.failedCount) || 0;
|
||||
@@ -527,6 +557,7 @@ export default function Accounts() {
|
||||
createdAccountId,
|
||||
recommendedModels,
|
||||
);
|
||||
await api.rebuildRoutes(false, false);
|
||||
seededRecommendedModels = true;
|
||||
} catch (seedErr: any) {
|
||||
toast.error(seedErr?.message || "连接已添加，但推荐模型补录失败");
|
||||
@@ -549,6 +580,52 @@ export default function Accounts() {
|
||||
`已补入 ${recommendedModels.length} 个推荐模型并重建路由`,
|
||||
);
|
||||
}
|
||||
const shouldAutoSyncModels =
|
||||
credentialMode === "apikey" &&
|
||||
!tokenForm.skipModelFetch &&
|
||||
createdAccountIds.length > 0;
|
||||
if (shouldAutoSyncModels) {
|
||||
if (
|
||||
!result?.batch &&
|
||||
verifyResult?.success &&
|
||||
Array.isArray(verifyResult.models) &&
|
||||
verifyResult.models.length > 0
|
||||
) {
|
||||
try {
|
||||
await api.addAccountAvailableModels(
|
||||
createdAccountIds[0],
|
||||
verifyResult.models,
|
||||
);
|
||||
await api.rebuildRoutes(false, false);
|
||||
toast.success("已自动同步模型并重建路由");
|
||||
} catch (syncErr: any) {
|
||||
toast.error(syncErr?.message || "连接已添加，但模型同步失败");
|
||||
}
|
||||
} else {
|
||||
const syncResults = await Promise.allSettled(
|
||||
createdAccountIds.map((accountId: number) =>
|
||||
api.checkModels(accountId),
|
||||
),
|
||||
);
|
||||
const syncedCount = syncResults.filter(
|
||||
(item) => item.status === "fulfilled",
|
||||
).length;
|
||||
const failedCount = syncResults.length - syncedCount;
|
||||
if (syncedCount > 0) {
|
||||
try {
|
||||
await api.rebuildRoutes(false, false);
|
||||
toast.success(
|
||||
`已自动同步 ${syncedCount} 个连接的模型并重建路由`,
|
||||
);
|
||||
} catch (rebuildErr: any) {
|
||||
toast.error(rebuildErr?.message || "模型已同步，但路由重建失败");
|
||||
}
|
||||
}
|
||||
if (failedCount > 0) {
|
||||
toast.error(`有 ${failedCount} 个连接同步模型失败`);
|
||||
}
|
||||
}
|
||||
}
|
||||
load(true);
|
||||
} catch (e: any) {
|
||||
toast.error(e.message || "添加失败");
|
||||
|
||||
+92
-134
@@ -1,16 +1,20 @@
|
||||
import React from 'react';
|
||||
import { useLocation } from 'react-router-dom';
|
||||
import { getSessionToken } from '../authSession.js';
|
||||
|
||||
export default function ComfyUI() {
|
||||
const [isDark, setIsDark] = React.useState(true);
|
||||
const [iframeSrc, setIframeSrc] = React.useState('');
|
||||
const [starting, setStarting] = React.useState(true);
|
||||
const [message, setMessage] = React.useState('正在启动 ComfyUI...');
|
||||
const [error, setError] = React.useState('');
|
||||
const location = useLocation();
|
||||
const audience = location.pathname.startsWith('/user/') ? 'user' : 'admin';
|
||||
const token = getSessionToken(audience, localStorage);
|
||||
const [available, setAvailable] = React.useState(false);
|
||||
const [message, setMessage] = React.useState('正在检查外部 ComfyUI...');
|
||||
const iframeRef = React.useRef<HTMLIFrameElement | null>(null);
|
||||
|
||||
const openFullscreen = () => {
|
||||
window.open('/api/comfyui/proxy/', '_blank', 'noopener,noreferrer');
|
||||
};
|
||||
|
||||
const reloadFrame = () => {
|
||||
if (iframeRef.current) {
|
||||
iframeRef.current.src = '/api/comfyui/proxy/';
|
||||
}
|
||||
};
|
||||
|
||||
React.useEffect(() => {
|
||||
const el = document.documentElement;
|
||||
@@ -23,121 +27,25 @@ export default function ComfyUI() {
|
||||
|
||||
React.useEffect(() => {
|
||||
let cancelled = false;
|
||||
const encodedToken = token ? encodeURIComponent(token) : '';
|
||||
const requestTimeoutMs = 35_000;
|
||||
|
||||
const fetchWithTimeout = async (input: RequestInfo | URL, init?: RequestInit) => {
|
||||
const controller = new AbortController();
|
||||
const timer = window.setTimeout(() => controller.abort(), requestTimeoutMs);
|
||||
try {
|
||||
return await fetch(input, { ...init, signal: controller.signal });
|
||||
} finally {
|
||||
window.clearTimeout(timer);
|
||||
}
|
||||
};
|
||||
|
||||
const resolveIframeFromPayload = (payload: any) => {
|
||||
const comfyui = payload?.comfyui || payload || {};
|
||||
const mode = String(comfyui?.mode || '').toLowerCase();
|
||||
const status = String(comfyui?.status || '').toLowerCase();
|
||||
const lastError = String(comfyui?.lastError || comfyui?.last_error || '').trim();
|
||||
|
||||
if (mode === 'container') {
|
||||
if (status === 'running') return `/comfy/?token=${encodedToken}`;
|
||||
if (status === 'unhealthy' || status === 'failed') {
|
||||
throw new Error(lastError || 'ComfyUI 容器启动失败');
|
||||
}
|
||||
return '';
|
||||
}
|
||||
|
||||
if (status && status !== 'disabled') {
|
||||
return '/api/comfyui/proxy/';
|
||||
}
|
||||
|
||||
if (lastError) {
|
||||
throw new Error(lastError);
|
||||
}
|
||||
|
||||
return '';
|
||||
};
|
||||
|
||||
if (!token) {
|
||||
setIframeSrc('');
|
||||
setStarting(false);
|
||||
setMessage('登录状态已过期，请重新登录。');
|
||||
setError('');
|
||||
void fetch('/api/comfyui/proxy/', { method: 'GET', signal: controller.signal })
|
||||
.then((response) => {
|
||||
if (cancelled) return;
|
||||
setAvailable(response.ok);
|
||||
setMessage(response.ok
|
||||
? ''
|
||||
: '外部 ComfyUI 不可达。请在 config.toml 中配置 comfyui.base_url，或使用工作流。');
|
||||
})
|
||||
.catch((err) => {
|
||||
if (cancelled || err?.name === 'AbortError') return;
|
||||
setAvailable(false);
|
||||
setMessage('外部 ComfyUI 不可达。请在 config.toml 中配置 comfyui.base_url，或使用工作流。');
|
||||
});
|
||||
return () => {
|
||||
cancelled = true;
|
||||
controller.abort();
|
||||
};
|
||||
}
|
||||
|
||||
setStarting(true);
|
||||
setIframeSrc('');
|
||||
setError('');
|
||||
setMessage('正在启动 ComfyUI...');
|
||||
|
||||
const start = async () => {
|
||||
try {
|
||||
const response = await fetchWithTimeout('/api/video-agent/start', {
|
||||
method: 'POST',
|
||||
headers: { Authorization: `Bearer ${token}` },
|
||||
});
|
||||
if (!cancelled) {
|
||||
if (!response.ok) {
|
||||
throw new Error(`启动接口返回 ${response.status}`);
|
||||
}
|
||||
const payload = await response.json().catch(() => null);
|
||||
const nextSrc = resolveIframeFromPayload(payload);
|
||||
if (nextSrc) {
|
||||
setIframeSrc(nextSrc);
|
||||
return;
|
||||
}
|
||||
}
|
||||
} catch (e: any) {
|
||||
if (!cancelled) {
|
||||
setMessage('ComfyUI 启动较慢，正在检查状态...');
|
||||
if (e?.name !== 'AbortError') setError(e?.message || '启动请求失败');
|
||||
}
|
||||
}
|
||||
|
||||
const deadline = Date.now() + 120_000;
|
||||
while (!cancelled && Date.now() < deadline) {
|
||||
try {
|
||||
const response = await fetchWithTimeout('/api/video-agent/status', {
|
||||
headers: { Authorization: `Bearer ${token}` },
|
||||
});
|
||||
if (response.ok) {
|
||||
const payload = await response.json().catch(() => null);
|
||||
const nextSrc = resolveIframeFromPayload(payload);
|
||||
if (nextSrc) {
|
||||
setIframeSrc(nextSrc);
|
||||
return;
|
||||
}
|
||||
setMessage('ComfyUI 仍在启动，等待容器就绪...');
|
||||
} else {
|
||||
setError(`状态接口返回 ${response.status}`);
|
||||
}
|
||||
} catch (e: any) {
|
||||
setError(e?.name === 'AbortError' ? '状态检查超时' : (e?.message || '状态检查失败'));
|
||||
}
|
||||
await new Promise((resolve) => window.setTimeout(resolve, 2000));
|
||||
}
|
||||
|
||||
if (!cancelled) {
|
||||
setStarting(false);
|
||||
setMessage('ComfyUI 启动超时');
|
||||
setError((current) => current || '请检查 /api/video-agent/status、Docker 容器日志和 4000 后端进程。');
|
||||
}
|
||||
};
|
||||
|
||||
void start().finally(() => {
|
||||
if (!cancelled) setStarting(false);
|
||||
});
|
||||
|
||||
return () => {
|
||||
cancelled = true;
|
||||
};
|
||||
}, [token]);
|
||||
}, []);
|
||||
|
||||
return (
|
||||
<div
|
||||
@@ -149,15 +57,23 @@ export default function ComfyUI() {
|
||||
background: isDark ? '#1a1a2e' : '#f5f5f5',
|
||||
}}
|
||||
>
|
||||
{iframeSrc ? (
|
||||
<div style={toolbarStyle}>
|
||||
<div style={toolbarTitleStyle}>外部 ComfyUI</div>
|
||||
<div style={toolbarActionsStyle}>
|
||||
<button type="button" onClick={reloadFrame} disabled={!available} style={toolbarButtonStyle}>
|
||||
刷新
|
||||
</button>
|
||||
<button type="button" onClick={openFullscreen} disabled={!available} style={primaryToolbarButtonStyle}>
|
||||
新窗口打开
|
||||
</button>
|
||||
</div>
|
||||
</div>
|
||||
{available ? (
|
||||
<iframe
|
||||
src={iframeSrc}
|
||||
style={{
|
||||
width: '100%',
|
||||
height: '100%',
|
||||
border: 'none',
|
||||
}}
|
||||
title="ComfyUI"
|
||||
ref={iframeRef}
|
||||
src="/api/comfyui/proxy/"
|
||||
style={{ width: '100%', height: '100%', border: 'none' }}
|
||||
title="External ComfyUI"
|
||||
allow="clipboard-read; clipboard-write"
|
||||
/>
|
||||
) : (
|
||||
@@ -168,16 +84,58 @@ export default function ComfyUI() {
|
||||
color: 'var(--color-text-muted)',
|
||||
fontSize: 14,
|
||||
textAlign: 'center',
|
||||
gap: 8,
|
||||
padding: 24,
|
||||
}}>
|
||||
<div>{starting ? message : message}</div>
|
||||
{error && (
|
||||
<div style={{ color: '#ef4444', fontSize: 12, maxWidth: 520 }}>
|
||||
{error}
|
||||
<div>
|
||||
<div style={{ fontWeight: 700, color: 'var(--color-text-primary)', marginBottom: 8 }}>
|
||||
内置 Docker ComfyUI 工作台已移除
|
||||
</div>
|
||||
<div style={{ maxWidth: 560, lineHeight: 1.7 }}>{message}</div>
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
const toolbarStyle: React.CSSProperties = {
|
||||
minHeight: 44,
|
||||
display: 'flex',
|
||||
alignItems: 'center',
|
||||
justifyContent: 'space-between',
|
||||
gap: 12,
|
||||
padding: '8px 12px',
|
||||
borderBottom: '1px solid var(--color-border-light)',
|
||||
background: 'var(--color-surface)',
|
||||
};
|
||||
|
||||
const toolbarTitleStyle: React.CSSProperties = {
|
||||
fontSize: 14,
|
||||
fontWeight: 800,
|
||||
color: 'var(--color-text-primary)',
|
||||
};
|
||||
|
||||
const toolbarActionsStyle: React.CSSProperties = {
|
||||
display: 'flex',
|
||||
alignItems: 'center',
|
||||
gap: 8,
|
||||
};
|
||||
|
||||
const toolbarButtonStyle: React.CSSProperties = {
|
||||
minHeight: 30,
|
||||
border: '1px solid var(--color-border)',
|
||||
borderRadius: 8,
|
||||
padding: '5px 10px',
|
||||
background: 'var(--color-bg)',
|
||||
color: 'var(--color-text-primary)',
|
||||
cursor: 'pointer',
|
||||
fontSize: 12,
|
||||
};
|
||||
|
||||
const primaryToolbarButtonStyle: React.CSSProperties = {
|
||||
...toolbarButtonStyle,
|
||||
borderColor: 'var(--color-primary)',
|
||||
background: 'var(--color-primary)',
|
||||
color: '#fff',
|
||||
fontWeight: 700,
|
||||
};
|
||||
|
||||
+860
-22
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,340 @@
|
||||
import React, { CSSProperties, useMemo, useState } from 'react';
|
||||
import { useLocation } from 'react-router-dom';
|
||||
import { fetchSessionResponse } from '../api.js';
|
||||
import type { SessionAudience } from '../authSession.js';
|
||||
|
||||
type Mode = 'audio' | 'image' | 'text';
|
||||
|
||||
const MODE_OPTIONS: Array<{ value: Mode; label: string; description: string }> = [
|
||||
{ value: 'audio', label: '音频驱动', description: '人物图 + 音频，适合口播、播报和虚拟主播。' },
|
||||
{ value: 'image', label: '图像驱动', description: '人物图 + 台词，交给上游数字人/视频模型生成。' },
|
||||
{ value: 'text', label: '文本生成', description: '只输入角色描述和台词，让模型生成数字人视频任务。' },
|
||||
];
|
||||
|
||||
const MODEL_PRESETS = [
|
||||
'longcat-video-avatar',
|
||||
'doubao-seedance-1-0-lite-i2v-250428',
|
||||
'doubao-seedance-1-0-lite-t2v-250428',
|
||||
'wan2-1-14b-i2v-250225',
|
||||
];
|
||||
|
||||
const VOICE_STYLE_OPTIONS = ['自然口播', '新闻播报', '带货讲解', '课程讲师', '双人对话'];
|
||||
|
||||
async function parseError(response: Response) {
|
||||
const text = await response.text().catch(() => '');
|
||||
if (!text) return `HTTP ${response.status}`;
|
||||
try {
|
||||
const value = JSON.parse(text);
|
||||
return value?.error?.message || value?.message || text;
|
||||
} catch {
|
||||
return text;
|
||||
}
|
||||
}
|
||||
|
||||
function extractTaskId(value: any): string {
|
||||
return value?.response?.id
|
||||
|| value?.response?.task_id
|
||||
|| value?.response?.data?.id
|
||||
|| value?.response?.data?.task_id
|
||||
|| value?.id
|
||||
|| value?.task_id
|
||||
|| '';
|
||||
}
|
||||
|
||||
function copyText(value: string) {
|
||||
if (!value) return;
|
||||
void navigator.clipboard?.writeText(value);
|
||||
}
|
||||
|
||||
export default function DigitalHuman() {
|
||||
const location = useLocation();
|
||||
const audience: SessionAudience = location.pathname.startsWith('/user/') ? 'user' : 'admin';
|
||||
const [mode, setMode] = useState<Mode>('audio');
|
||||
const [model, setModel] = useState(MODEL_PRESETS[0]);
|
||||
const [prompt, setPrompt] = useState('一位专业中文数字人主播，坐在简洁演播室中，面向镜头自然讲解 BoosAPI 的模型路由和工作流能力。');
|
||||
const [script, setScript] = useState('大家好，欢迎使用 BoosAPI。今天我们演示如何用统一模型路由生成数字人口播视频。');
|
||||
const [imageUrl, setImageUrl] = useState('');
|
||||
const [audioUrl, setAudioUrl] = useState('');
|
||||
const [voiceStyle, setVoiceStyle] = useState(VOICE_STYLE_OPTIONS[0]);
|
||||
const [resolution, setResolution] = useState('720p');
|
||||
const [ratio, setRatio] = useState('9:16');
|
||||
const [duration, setDuration] = useState(8);
|
||||
const [refImageIndex, setRefImageIndex] = useState(0);
|
||||
const [audioCfg, setAudioCfg] = useState(4.5);
|
||||
const [segments, setSegments] = useState(false);
|
||||
const [multiPerson, setMultiPerson] = useState(false);
|
||||
const [submitting, setSubmitting] = useState(false);
|
||||
const [error, setError] = useState('');
|
||||
const [result, setResult] = useState<any>(null);
|
||||
|
||||
const bodyPreview = useMemo(() => ({
|
||||
model: model.trim(),
|
||||
prompt: [
|
||||
prompt.trim(),
|
||||
script.trim() ? `台词：${script.trim()}` : '',
|
||||
`风格：${voiceStyle}`,
|
||||
].filter(Boolean).join('\n'),
|
||||
mode,
|
||||
ratio,
|
||||
resolution,
|
||||
duration,
|
||||
audio_url: mode === 'audio' ? audioUrl.trim() : undefined,
|
||||
image_url: mode !== 'text' ? imageUrl.trim() : undefined,
|
||||
audio_cfg: audioCfg,
|
||||
ref_img_index: refImageIndex,
|
||||
segmented: segments,
|
||||
multi_person: multiPerson,
|
||||
}), [audioCfg, audioUrl, duration, imageUrl, mode, model, multiPerson, prompt, ratio, refImageIndex, resolution, script, segments, voiceStyle]);
|
||||
|
||||
const submit = async () => {
|
||||
setSubmitting(true);
|
||||
setError('');
|
||||
setResult(null);
|
||||
try {
|
||||
const response = await fetchSessionResponse(audience, '/api/digital-human/generate', {
|
||||
method: 'POST',
|
||||
body: JSON.stringify(bodyPreview),
|
||||
timeoutMs: 300_000,
|
||||
});
|
||||
if (!response.ok) throw new Error(await parseError(response));
|
||||
setResult(await response.json());
|
||||
} catch (err: any) {
|
||||
setError(err?.message || '数字人任务创建失败');
|
||||
} finally {
|
||||
setSubmitting(false);
|
||||
}
|
||||
};
|
||||
|
||||
const taskId = extractTaskId(result);
|
||||
|
||||
return (
|
||||
<div className="page">
|
||||
<div className="page-header">
|
||||
<div>
|
||||
<h1 className="page-title">数字人</h1>
|
||||
<p className="page-subtitle">参考 LongCat-Video Avatar 的任务形态：人物图、音频、台词、分段和多人数字人编排，底层走 BoosAPI 视频模型路由。</p>
|
||||
</div>
|
||||
<div className="page-actions">
|
||||
<button type="button" className="btn btn-primary" onClick={submit} disabled={submitting}>
|
||||
{submitting ? '创建中...' : '创建数字人任务'}
|
||||
</button>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
{error && <div className="alert alert-error" style={{ marginBottom: 12 }}>{error}</div>}
|
||||
|
||||
<div style={layoutStyle}>
|
||||
<section style={panelStyle}>
|
||||
<div style={modeGridStyle}>
|
||||
{MODE_OPTIONS.map((item) => (
|
||||
<button
|
||||
key={item.value}
|
||||
type="button"
|
||||
onClick={() => setMode(item.value)}
|
||||
style={mode === item.value ? activeModeStyle : modeStyle}
|
||||
>
|
||||
<strong>{item.label}</strong>
|
||||
<span>{item.description}</span>
|
||||
</button>
|
||||
))}
|
||||
</div>
|
||||
|
||||
<div style={formGridStyle}>
|
||||
<label style={fieldStyle}>
|
||||
<span>模型</span>
|
||||
<input value={model} onChange={(event) => setModel(event.target.value)} list="digital-human-models" style={inputStyle} />
|
||||
<datalist id="digital-human-models">
|
||||
{MODEL_PRESETS.map((item) => <option key={item} value={item} />)}
|
||||
</datalist>
|
||||
</label>
|
||||
<label style={fieldStyle}>
|
||||
<span>画幅</span>
|
||||
<select value={ratio} onChange={(event) => setRatio(event.target.value)} style={inputStyle}>
|
||||
<option value="9:16">9:16 竖屏</option>
|
||||
<option value="16:9">16:9 横屏</option>
|
||||
<option value="1:1">1:1 方屏</option>
|
||||
</select>
|
||||
</label>
|
||||
<label style={fieldStyle}>
|
||||
<span>分辨率</span>
|
||||
<select value={resolution} onChange={(event) => setResolution(event.target.value)} style={inputStyle}>
|
||||
<option value="720p">720p</option>
|
||||
<option value="1080p">1080p</option>
|
||||
<option value="480p">480p</option>
|
||||
</select>
|
||||
</label>
|
||||
<label style={fieldStyle}>
|
||||
<span>时长（秒）</span>
|
||||
<input type="number" min={1} max={120} value={duration} onChange={(event) => setDuration(Number(event.target.value) || 1)} style={inputStyle} />
|
||||
</label>
|
||||
</div>
|
||||
|
||||
<label style={fieldStyle}>
|
||||
<span>人物/场景描述</span>
|
||||
<textarea value={prompt} onChange={(event) => setPrompt(event.target.value)} rows={4} style={textareaStyle} />
|
||||
</label>
|
||||
<label style={fieldStyle}>
|
||||
<span>台词/脚本</span>
|
||||
<textarea value={script} onChange={(event) => setScript(event.target.value)} rows={5} style={textareaStyle} />
|
||||
</label>
|
||||
|
||||
<div style={formGridStyle}>
|
||||
{mode !== 'text' && (
|
||||
<label style={fieldStyle}>
|
||||
<span>人物图 URL</span>
|
||||
<input value={imageUrl} onChange={(event) => setImageUrl(event.target.value)} placeholder="https://..." style={inputStyle} />
|
||||
</label>
|
||||
)}
|
||||
{mode === 'audio' && (
|
||||
<label style={fieldStyle}>
|
||||
<span>音频 URL</span>
|
||||
<input value={audioUrl} onChange={(event) => setAudioUrl(event.target.value)} placeholder="https://..." style={inputStyle} />
|
||||
</label>
|
||||
)}
|
||||
<label style={fieldStyle}>
|
||||
<span>口播风格</span>
|
||||
<select value={voiceStyle} onChange={(event) => setVoiceStyle(event.target.value)} style={inputStyle}>
|
||||
{VOICE_STYLE_OPTIONS.map((item) => <option key={item} value={item}>{item}</option>)}
|
||||
</select>
|
||||
</label>
|
||||
<label style={fieldStyle}>
|
||||
<span>ref_img_index</span>
|
||||
<input type="number" min={0} value={refImageIndex} onChange={(event) => setRefImageIndex(Number(event.target.value) || 0)} style={inputStyle} />
|
||||
</label>
|
||||
<label style={fieldStyle}>
|
||||
<span>audio CFG</span>
|
||||
<input type="number" min={0} step={0.1} value={audioCfg} onChange={(event) => setAudioCfg(Number(event.target.value) || 0)} style={inputStyle} />
|
||||
</label>
|
||||
</div>
|
||||
|
||||
<div style={toggleRowStyle}>
|
||||
<label><input type="checkbox" checked={segments} onChange={(event) => setSegments(event.target.checked)} /> 分段长视频</label>
|
||||
<label><input type="checkbox" checked={multiPerson} onChange={(event) => setMultiPerson(event.target.checked)} /> 多人数字人</label>
|
||||
</div>
|
||||
</section>
|
||||
|
||||
<aside style={panelStyle}>
|
||||
<h2 style={sideTitleStyle}>任务请求</h2>
|
||||
<pre style={codeStyle}>{JSON.stringify(bodyPreview, null, 2)}</pre>
|
||||
<button type="button" className="btn btn-ghost" onClick={() => copyText(JSON.stringify(bodyPreview, null, 2))}>复制请求</button>
|
||||
|
||||
<h2 style={{ ...sideTitleStyle, marginTop: 18 }}>返回结果</h2>
|
||||
{result ? (
|
||||
<>
|
||||
{taskId && <div style={taskIdStyle}>任务 ID：{taskId}</div>}
|
||||
<pre style={codeStyle}>{JSON.stringify(result.response || result, null, 2)}</pre>
|
||||
<button type="button" className="btn btn-ghost" onClick={() => copyText(JSON.stringify(result.response || result, null, 2))}>复制结果</button>
|
||||
</>
|
||||
) : (
|
||||
<div className="empty-state" style={{ minHeight: 160 }}>
|
||||
<div className="empty-state-title">暂无任务</div>
|
||||
<div className="empty-state-desc">创建后会显示上游返回，可用任务 ID 到模型站点或后续查询接口查看状态。</div>
|
||||
</div>
|
||||
)}
|
||||
</aside>
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
const layoutStyle: CSSProperties = {
|
||||
display: 'grid',
|
||||
gridTemplateColumns: 'minmax(0, 1.25fr) minmax(320px, 0.75fr)',
|
||||
gap: 16,
|
||||
};
|
||||
|
||||
const panelStyle: CSSProperties = {
|
||||
border: '1px solid var(--color-border)',
|
||||
background: 'var(--color-surface)',
|
||||
borderRadius: 8,
|
||||
padding: 16,
|
||||
};
|
||||
|
||||
const modeGridStyle: CSSProperties = {
|
||||
display: 'grid',
|
||||
gridTemplateColumns: 'repeat(auto-fit, minmax(180px, 1fr))',
|
||||
gap: 10,
|
||||
marginBottom: 16,
|
||||
};
|
||||
|
||||
const modeStyle: CSSProperties = {
|
||||
textAlign: 'left',
|
||||
border: '1px solid var(--color-border)',
|
||||
background: 'var(--color-bg)',
|
||||
borderRadius: 8,
|
||||
padding: 12,
|
||||
display: 'flex',
|
||||
flexDirection: 'column',
|
||||
gap: 6,
|
||||
color: 'var(--color-text-primary)',
|
||||
};
|
||||
|
||||
const activeModeStyle: CSSProperties = {
|
||||
...modeStyle,
|
||||
borderColor: 'var(--color-primary)',
|
||||
boxShadow: '0 0 0 2px color-mix(in srgb, var(--color-primary) 15%, transparent)',
|
||||
};
|
||||
|
||||
const formGridStyle: CSSProperties = {
|
||||
display: 'grid',
|
||||
gridTemplateColumns: 'repeat(auto-fit, minmax(180px, 1fr))',
|
||||
gap: 12,
|
||||
marginBottom: 12,
|
||||
};
|
||||
|
||||
const fieldStyle: CSSProperties = {
|
||||
display: 'flex',
|
||||
flexDirection: 'column',
|
||||
gap: 6,
|
||||
fontSize: 12,
|
||||
color: 'var(--color-text-secondary)',
|
||||
};
|
||||
|
||||
const inputStyle: CSSProperties = {
|
||||
width: '100%',
|
||||
border: '1px solid var(--color-border)',
|
||||
background: 'var(--color-bg)',
|
||||
color: 'var(--color-text-primary)',
|
||||
borderRadius: 6,
|
||||
padding: '10px 12px',
|
||||
fontSize: 13,
|
||||
};
|
||||
|
||||
const textareaStyle: CSSProperties = {
|
||||
...inputStyle,
|
||||
minHeight: 110,
|
||||
resize: 'vertical',
|
||||
};
|
||||
|
||||
const toggleRowStyle: CSSProperties = {
|
||||
display: 'flex',
|
||||
gap: 16,
|
||||
flexWrap: 'wrap',
|
||||
fontSize: 13,
|
||||
color: 'var(--color-text-secondary)',
|
||||
};
|
||||
|
||||
const sideTitleStyle: CSSProperties = {
|
||||
margin: '0 0 10px',
|
||||
fontSize: 14,
|
||||
};
|
||||
|
||||
const codeStyle: CSSProperties = {
|
||||
maxHeight: 280,
|
||||
overflow: 'auto',
|
||||
background: 'var(--color-bg)',
|
||||
border: '1px solid var(--color-border)',
|
||||
borderRadius: 6,
|
||||
padding: 12,
|
||||
fontSize: 12,
|
||||
whiteSpace: 'pre-wrap',
|
||||
};
|
||||
|
||||
const taskIdStyle: CSSProperties = {
|
||||
marginBottom: 8,
|
||||
padding: '8px 10px',
|
||||
borderRadius: 6,
|
||||
background: 'color-mix(in srgb, var(--color-primary) 10%, transparent)',
|
||||
color: 'var(--color-text-primary)',
|
||||
fontSize: 12,
|
||||
};
|
||||
+92
-27
@@ -1,6 +1,6 @@
|
||||
import React, { useCallback, useEffect, useMemo, useRef, useState } from 'react';
|
||||
import { useLocation } from 'react-router-dom';
|
||||
import { api } from '../api.js';
|
||||
import { api, userApi } from '../api.js';
|
||||
import { BrandGlyph, getBrand, hashColor, BrandIcon, type BrandInfo } from '../components/BrandIcon.js';
|
||||
import SiteBadgeLink from '../components/SiteBadgeLink.js';
|
||||
import { useToast } from '../components/Toast.js';
|
||||
@@ -125,6 +125,24 @@ function renderGroupPricingValue(pricing: ModelGroupPricing): string {
|
||||
return `${pricing.perCallTotal ?? 0} USD / call`;
|
||||
}
|
||||
|
||||
async function copyTextToClipboard(text: string) {
|
||||
if (navigator.clipboard?.writeText) {
|
||||
await navigator.clipboard.writeText(text);
|
||||
return;
|
||||
}
|
||||
|
||||
const textarea = document.createElement('textarea');
|
||||
textarea.value = text;
|
||||
textarea.style.position = 'fixed';
|
||||
textarea.style.opacity = '0';
|
||||
textarea.style.left = '-9999px';
|
||||
document.body.appendChild(textarea);
|
||||
textarea.focus();
|
||||
textarea.select();
|
||||
document.execCommand('copy');
|
||||
document.body.removeChild(textarea);
|
||||
}
|
||||
|
||||
const PAGE_SIZES = [10, 20, 50];
|
||||
|
||||
function compareModels(a: ModelRow, b: ModelRow, sortBy: SortColumn, sortDir: 'asc' | 'desc'): number {
|
||||
@@ -153,6 +171,9 @@ function compareModels(a: ModelRow, b: ModelRow, sortBy: SortColumn, sortDir: 'a
|
||||
/* ---- component ---- */
|
||||
export default function Models() {
|
||||
const toast = useToast();
|
||||
const location = useLocation();
|
||||
const isUserView = location.pathname.startsWith('/user/');
|
||||
const marketplaceApi = isUserView ? userApi : api;
|
||||
const [data, setData] = useState<ModelsMarketplaceResponse>({ models: [] });
|
||||
const [loading, setLoading] = useState(true);
|
||||
const [search, setSearch] = useState('');
|
||||
@@ -172,7 +193,6 @@ export default function Models() {
|
||||
const filterPanelPresence = useAnimatedVisibility(!isMobile && !filterCollapsed, 220);
|
||||
const latestPrimaryRequestRef = useRef(0);
|
||||
const latestMetadataRequestRef = useRef(0);
|
||||
const location = useLocation();
|
||||
const siteIdByName = useMemo(() => {
|
||||
const index = new Map<string, number>();
|
||||
for (const model of data.models) {
|
||||
@@ -192,7 +212,7 @@ export default function Models() {
|
||||
setMetadataHydrating(false);
|
||||
setLoading(true);
|
||||
try {
|
||||
const res = await api.getModelsMarketplace({
|
||||
const res = await marketplaceApi.getModelsMarketplace({
|
||||
refresh,
|
||||
includePricing: false,
|
||||
});
|
||||
@@ -216,7 +236,7 @@ export default function Models() {
|
||||
setLoading(false);
|
||||
}
|
||||
}
|
||||
}, [toast]);
|
||||
}, [marketplaceApi, toast]);
|
||||
|
||||
const hydrateMarketplaceMetadata = useCallback(async (baseModels: ModelRow[]) => {
|
||||
if (!shouldHydrateMarketplaceMetadata(baseModels)) return;
|
||||
@@ -225,7 +245,7 @@ export default function Models() {
|
||||
const baseRequestId = latestPrimaryRequestRef.current;
|
||||
setMetadataHydrating(true);
|
||||
try {
|
||||
const res = await api.getModelsMarketplace({
|
||||
const res = await marketplaceApi.getModelsMarketplace({
|
||||
includePricing: true,
|
||||
});
|
||||
if (metadataRequestId !== latestMetadataRequestRef.current) return;
|
||||
@@ -244,7 +264,7 @@ export default function Models() {
|
||||
setMetadataHydrating(false);
|
||||
}
|
||||
}
|
||||
}, []);
|
||||
}, [marketplaceApi]);
|
||||
|
||||
useEffect(() => {
|
||||
if (!isMobile) return;
|
||||
@@ -311,6 +331,7 @@ export default function Models() {
|
||||
|
||||
/* ---- derived: site list ---- */
|
||||
const siteMap = useMemo(() => {
|
||||
if (isUserView) return [];
|
||||
const m = new Map<string, number>();
|
||||
for (const model of data.models) {
|
||||
for (const a of model.accounts) {
|
||||
@@ -318,7 +339,7 @@ export default function Models() {
|
||||
}
|
||||
}
|
||||
return [...m.entries()].sort((a, b) => b[1] - a[1]);
|
||||
}, [data.models]);
|
||||
}, [data.models, isUserView]);
|
||||
|
||||
/* ---- filtered ---- */
|
||||
const filteredModels = useMemo(() => {
|
||||
@@ -366,8 +387,8 @@ export default function Models() {
|
||||
};
|
||||
}) : filteredModels;
|
||||
|
||||
return [...scopedModels].sort((a, b) => compareModels(a, b, sortBy, sortDir));
|
||||
}, [filteredModels, activeSite, sortBy, sortDir]);
|
||||
return [...scopedModels].sort((a, b) => compareModels(a, b, isUserView ? 'name' : sortBy, isUserView ? 'asc' : sortDir));
|
||||
}, [filteredModels, activeSite, sortBy, sortDir, isUserView]);
|
||||
|
||||
/* ---- pagination ---- */
|
||||
const totalPages = Math.max(1, Math.ceil(detailModels.length / pageSize));
|
||||
@@ -395,10 +416,15 @@ export default function Models() {
|
||||
: null;
|
||||
|
||||
/* ---- copy ---- */
|
||||
const copyName = (name: string) => {
|
||||
navigator.clipboard.writeText(name).catch(() => { });
|
||||
const copyName = async (name: string) => {
|
||||
try {
|
||||
await copyTextToClipboard(name);
|
||||
setCopied(name);
|
||||
toast.success(tr('已复制模型名'));
|
||||
setTimeout(() => setCopied(null), 1500);
|
||||
} catch (error: any) {
|
||||
toast.error(error?.message || tr('复制模型名失败'));
|
||||
}
|
||||
};
|
||||
|
||||
const filterControls = (
|
||||
@@ -441,6 +467,7 @@ export default function Models() {
|
||||
)}
|
||||
</div>
|
||||
|
||||
{!isUserView && (
|
||||
<div className="filter-panel-section">
|
||||
<div className="filter-panel-title">
|
||||
{tr('供应商')}
|
||||
@@ -460,8 +487,9 @@ export default function Models() {
|
||||
</div>
|
||||
))}
|
||||
</div>
|
||||
)}
|
||||
|
||||
<div className="filter-panel-section">
|
||||
{!isUserView && <div className="filter-panel-section">
|
||||
<div className="filter-panel-title">{tr('排序方式')}</div>
|
||||
{[
|
||||
{ key: 'accountCount' as SortColumn, label: tr('账号数') },
|
||||
@@ -490,7 +518,7 @@ export default function Models() {
|
||||
)}
|
||||
</div>
|
||||
))}
|
||||
</div>
|
||||
</div>}
|
||||
</>
|
||||
);
|
||||
|
||||
@@ -626,6 +654,7 @@ export default function Models() {
|
||||
/>
|
||||
</div>
|
||||
{/* Quick stats */}
|
||||
{!isUserView && (
|
||||
<div style={{ display: 'flex', gap: 16, fontSize: 12, color: 'var(--color-text-muted)', alignItems: 'center' }}>
|
||||
<span data-tooltip={tr('所有模型 accountCount 累计值，同一账号在多个模型中会重复计数')}>
|
||||
{tr('覆盖档位')} <b style={{ color: 'var(--color-text-primary)' }}>{totalCoverageSlots}</b>
|
||||
@@ -635,6 +664,7 @@ export default function Models() {
|
||||
</span>
|
||||
<span>{tr('平均延迟')} <b style={{ color: getMetricColor(avgLatency) }}>{formatLatency(avgLatency)}</b></span>
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
|
||||
{/* Empty */}
|
||||
@@ -660,14 +690,14 @@ export default function Models() {
|
||||
<div className="model-card-info">
|
||||
<div className="model-card-name">{m.name}</div>
|
||||
<div className="model-card-meta">
|
||||
<span>
|
||||
{!isUserView && <span>
|
||||
<svg width="12" height="12" fill="none" viewBox="0 0 24 24" stroke="currentColor"><path strokeLinecap="round" strokeLinejoin="round" strokeWidth={2} d="M17 20h5v-2a3 3 0 00-5.356-1.857M17 20H7m10 0v-2c0-.656-.126-1.283-.356-1.857M7 20H2v-2a3 3 0 015.356-1.857M7 20v-2c0-.656.126-1.283.356-1.857m0 0a5.002 5.002 0 019.288 0M15 7a3 3 0 11-6 0 3 3 0 016 0z" /></svg>
|
||||
{m.accountCount} {tr('个账号')}
|
||||
</span>
|
||||
<span>
|
||||
</span>}
|
||||
{!isUserView && <span>
|
||||
<svg width="12" height="12" fill="none" viewBox="0 0 24 24" stroke="currentColor"><path strokeLinecap="round" strokeLinejoin="round" strokeWidth={2} d="M15 7a2 2 0 012 2m4 0a6 6 0 01-7.743 5.743L11 17H9v2H7v2H4a1 1 0 01-1-1v-2.586a1 1 0 01.293-.707l5.964-5.964A6 6 0 1121 9z" /></svg>
|
||||
{m.tokenCount} {tr('令牌')}
|
||||
</span>
|
||||
</span>}
|
||||
<span
|
||||
className={`badge ${getLatencyBadgeClass(m.avgLatency)}`}
|
||||
style={{ fontVariantNumeric: 'tabular-nums' }}
|
||||
@@ -685,7 +715,16 @@ export default function Models() {
|
||||
</div>
|
||||
</div>
|
||||
<div className="model-card-actions" onClick={e => e.stopPropagation()}>
|
||||
<button className="model-card-action-btn" data-tooltip={tr('复制模型名')} aria-label={tr('复制模型名')} onClick={() => copyName(m.name)}>
|
||||
<button
|
||||
type="button"
|
||||
className="model-card-action-btn"
|
||||
data-tooltip={tr('复制模型名')}
|
||||
aria-label={tr('复制模型名')}
|
||||
onClick={(event) => {
|
||||
event?.stopPropagation?.();
|
||||
void copyName(m.name);
|
||||
}}
|
||||
>
|
||||
{copied === m.name ? (
|
||||
<svg width="14" height="14" fill="none" viewBox="0 0 24 24" stroke="var(--color-success)"><path strokeLinecap="round" strokeLinejoin="round" strokeWidth={2} d="M5 13l4 4L19 7" /></svg>
|
||||
) : (
|
||||
@@ -709,7 +748,7 @@ export default function Models() {
|
||||
{getBrand(m.name) && (
|
||||
<span className="model-tag model-tag-purple">{getBrand(m.name)!.name}</span>
|
||||
)}
|
||||
{m.accounts.map(a => a.site).filter((v, i, arr) => arr.indexOf(v) === i).map(site => (
|
||||
{!isUserView && m.accounts.map(a => a.site).filter((v, i, arr) => arr.indexOf(v) === i).map(site => (
|
||||
<span key={site} className="model-tag model-tag-blue">{site}</span>
|
||||
))}
|
||||
{m.successRate != null && m.successRate >= 90 && (
|
||||
@@ -750,6 +789,7 @@ export default function Models() {
|
||||
</div>
|
||||
</div>
|
||||
|
||||
{!isUserView && (
|
||||
<div className="card" style={{ padding: 10 }}>
|
||||
<div style={{ fontSize: 12, fontWeight: 600, marginBottom: 6 }}>{tr('分组计费')}</div>
|
||||
{m.pricingSources.length > 0 ? (
|
||||
@@ -760,7 +800,7 @@ export default function Models() {
|
||||
style={{ border: '1px solid var(--color-border-light)', borderRadius: 8, padding: 8 }}
|
||||
>
|
||||
<div style={{ fontSize: 12, marginBottom: 6 }}>
|
||||
<SiteBadgeLink siteId={source.siteId} siteName={source.siteName} badgeStyle={{ fontSize: 11 }} /> · {source.username || `ID:${source.accountId}`}
|
||||
{isUserView ? tr('隐藏账号详情') : <><SiteBadgeLink siteId={source.siteId} siteName={source.siteName} badgeStyle={{ fontSize: 11 }} /> · {source.username || `ID:${source.accountId}`}</>}
|
||||
</div>
|
||||
<div style={{ display: 'flex', gap: 6, flexWrap: 'wrap' }}>
|
||||
{Object.entries(source.groupPricing).map(([group, pricing]) => (
|
||||
@@ -776,9 +816,21 @@ export default function Models() {
|
||||
<span className="badge badge-muted">{metadataHydrating ? tr('正在加载价格元数据...') : tr('暂无价格元数据')}</span>
|
||||
)}
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
|
||||
{isMobile ? (
|
||||
{isUserView ? (
|
||||
<div style={{ display: 'grid', gap: 8 }}>
|
||||
<div style={{ fontSize: 12, fontWeight: 600, color: 'var(--color-text-secondary)' }}>{tr('品牌')}</div>
|
||||
<div className="card" style={{ padding: 10 }}>
|
||||
{getBrand(m.name) ? (
|
||||
<span className="model-tag model-tag-purple">{getBrand(m.name)!.name}</span>
|
||||
) : (
|
||||
<span className="badge badge-muted">{tr('未识别品牌')}</span>
|
||||
)}
|
||||
</div>
|
||||
</div>
|
||||
) : isMobile ? (
|
||||
<div style={{ display: 'grid', gap: 8 }}>
|
||||
<div style={{ fontSize: 12, fontWeight: 600, color: 'var(--color-text-secondary)' }}>{tr('账号明细')}</div>
|
||||
{m.accounts.map((a) => (
|
||||
@@ -828,8 +880,8 @@ export default function Models() {
|
||||
<tbody>
|
||||
{m.accounts.map(a => (
|
||||
<tr key={a.id}>
|
||||
<td><SiteBadgeLink siteId={siteIdByName.get(a.site)} siteName={a.site} badgeClassName="badge badge-info" badgeStyle={{ fontSize: 11 }} /></td>
|
||||
<td style={{ fontSize: 12 }}>{a.username || `ID:${a.id}`}</td>
|
||||
<td>{isUserView ? <span className="badge badge-muted">{tr('隐藏')}</span> : <SiteBadgeLink siteId={siteIdByName.get(a.site)} siteName={a.site} badgeClassName="badge badge-info" badgeStyle={{ fontSize: 11 }} />}</td>
|
||||
<td style={{ fontSize: 12 }}>{isUserView ? tr('隐藏') : (a.username || `ID:${a.id}`)}</td>
|
||||
<td style={{ display: 'flex', gap: 4, flexWrap: 'wrap' }}>
|
||||
{a.tokens.length > 0 ? a.tokens.map(t => (
|
||||
<span key={t.id} className={`badge ${t.isDefault ? 'badge-success' : 'badge-muted'}`} style={{ fontSize: 11 }}>{t.name}</span>
|
||||
@@ -864,12 +916,16 @@ export default function Models() {
|
||||
<th style={{ cursor: 'pointer' }} onClick={() => { setSortBy('name'); setSortDir(d => d === 'asc' ? 'desc' : 'asc'); }}>
|
||||
{tr('模型名称')} {sortBy === 'name' ? (sortDir === 'desc' ? '↓' : '↑') : ''}
|
||||
</th>
|
||||
{!isUserView && (
|
||||
<th style={{ cursor: 'pointer' }} onClick={() => { setSortBy('accountCount'); setSortDir(d => d === 'asc' ? 'desc' : 'asc'); }}>
|
||||
{tr('账号数')} {sortBy === 'accountCount' ? (sortDir === 'desc' ? '↓' : '↑') : ''}
|
||||
</th>
|
||||
)}
|
||||
{!isUserView && (
|
||||
<th style={{ cursor: 'pointer' }} onClick={() => { setSortBy('tokenCount'); setSortDir(d => d === 'asc' ? 'desc' : 'asc'); }}>
|
||||
{tr('令牌数')} {sortBy === 'tokenCount' ? (sortDir === 'desc' ? '↓' : '↑') : ''}
|
||||
</th>
|
||||
)}
|
||||
<th style={{ cursor: 'pointer' }} onClick={() => { setSortBy('avgLatency'); setSortDir(d => d === 'asc' ? 'desc' : 'asc'); }}>
|
||||
{tr('延迟')} {sortBy === 'avgLatency' ? (sortDir === 'desc' ? '↓' : '↑') : ''}
|
||||
</th>
|
||||
@@ -893,8 +949,8 @@ export default function Models() {
|
||||
{m.name}
|
||||
</code>
|
||||
</td>
|
||||
<td><span className="badge badge-info">{m.accountCount}</span></td>
|
||||
<td><span className="badge badge-muted">{m.tokenCount}</span></td>
|
||||
{!isUserView && <td><span className="badge badge-info">{m.accountCount}</span></td>}
|
||||
{!isUserView && <td><span className="badge badge-muted">{m.tokenCount}</span></td>}
|
||||
<td>
|
||||
<span
|
||||
className={`badge ${getLatencyBadgeClass(m.avgLatency)}`}
|
||||
@@ -912,7 +968,16 @@ export default function Models() {
|
||||
</span>
|
||||
</td>
|
||||
<td onClick={e => e.stopPropagation()}>
|
||||
<button className="model-card-action-btn" data-tooltip={tr('复制')} aria-label={tr('复制')} onClick={() => copyName(m.name)}>
|
||||
<button
|
||||
type="button"
|
||||
className="model-card-action-btn"
|
||||
data-tooltip={tr('复制')}
|
||||
aria-label={tr('复制')}
|
||||
onClick={(event) => {
|
||||
event?.stopPropagation?.();
|
||||
void copyName(m.name);
|
||||
}}
|
||||
>
|
||||
{copied === m.name ? (
|
||||
<svg width="12" height="12" fill="none" viewBox="0 0 24 24" stroke="var(--color-success)"><path strokeLinecap="round" strokeLinejoin="round" strokeWidth={2} d="M5 13l4 4L19 7" /></svg>
|
||||
) : (
|
||||
@@ -923,7 +988,7 @@ export default function Models() {
|
||||
</tr>
|
||||
{isExpanded ? (
|
||||
<tr className="log-detail-row">
|
||||
<td colSpan={7} style={{ padding: 0 }}>
|
||||
<td colSpan={isUserView ? 5 : 7} style={{ padding: 0 }}>
|
||||
<div className="anim-collapse is-open">
|
||||
<div className="anim-collapse-inner">
|
||||
<div style={{ padding: '12px 16px 12px 54px' }}>
|
||||
|
||||
+390
-13
@@ -13,6 +13,36 @@ type AgentImage = {
|
||||
};
|
||||
|
||||
type PictureMode = 'generate' | 'edit';
|
||||
type PictureJobStatus = 'pending' | 'running' | 'succeeded' | 'failed' | 'cancelled';
|
||||
|
||||
type PictureJob = {
|
||||
id: string;
|
||||
status: PictureJobStatus | string;
|
||||
model?: string | null;
|
||||
message?: string | null;
|
||||
error?: string | null;
|
||||
progress?: number;
|
||||
progressTotal?: number;
|
||||
result?: {
|
||||
prompt?: string;
|
||||
model?: string;
|
||||
mode?: PictureMode | string;
|
||||
response?: any;
|
||||
images?: any[];
|
||||
};
|
||||
createdAt?: string;
|
||||
updatedAt?: string;
|
||||
finishedAt?: string | null;
|
||||
};
|
||||
|
||||
function staleRunningHint(job: PictureJob) {
|
||||
if (job.status !== 'pending' && job.status !== 'running') return '';
|
||||
const updatedAt = parseServerTime(job.updatedAt || job.createdAt);
|
||||
if (!Number.isFinite(updatedAt)) return '';
|
||||
const ageMs = Date.now() - updatedAt;
|
||||
if (ageMs < 180_000) return '';
|
||||
return '图片模型响应时间较长，任务仍在后台等待。大型提示词或高峰期通常需要 3-5 分钟。';
|
||||
}
|
||||
|
||||
type PictureAgentStatus = {
|
||||
success?: boolean;
|
||||
@@ -47,6 +77,7 @@ type GalleryItem = {
|
||||
|
||||
const DEFAULT_PROMPT = 'A precise editorial product photo of a translucent glass robot painting a small neon city, sharp focus, clean background';
|
||||
const TEMPLATE_GALLERY_JSON_URL = 'https://weiyyin.github.io/linux-do-gpt-image/linux_do_2036927_prompts_images.filtered.json';
|
||||
const TEMPLATE_LIMIT = 144;
|
||||
const QUALITY_OPTIONS = ['auto', 'low', 'medium', 'high'];
|
||||
const ASPECT_OPTIONS = [
|
||||
{ value: '1:1', label: '1:1 方图', width: 1024, height: 1024 },
|
||||
@@ -84,7 +115,13 @@ function makeId() {
|
||||
}
|
||||
|
||||
function extractImageUrls(payload: any): string[] {
|
||||
const data = Array.isArray(payload?.data) ? payload.data : [];
|
||||
const data = Array.isArray(payload?.data)
|
||||
? payload.data
|
||||
: Array.isArray(payload?.response?.data)
|
||||
? payload.response.data
|
||||
: Array.isArray(payload?.images)
|
||||
? payload.images
|
||||
: [];
|
||||
return data
|
||||
.map((item: any) => {
|
||||
if (typeof item?.url === 'string' && item.url.trim()) return item.url.trim();
|
||||
@@ -112,6 +149,28 @@ function getFilename(image: AgentImage) {
|
||||
return `shengtu-${safeModel}-${new Date(image.createdAt).toISOString().slice(0, 19).replace(/[:T]/g, '-')}.png`;
|
||||
}
|
||||
|
||||
function parseServerTime(value?: string | null): number {
|
||||
if (!value) return Number.NaN;
|
||||
const trimmed = value.trim();
|
||||
if (!trimmed) return Number.NaN;
|
||||
const normalized = trimmed.includes('T') ? trimmed : trimmed.replace(' ', 'T');
|
||||
const withTimezone = /(?:Z|[+-]\d{2}:?\d{2})$/i.test(normalized)
|
||||
? normalized
|
||||
: `${normalized}Z`;
|
||||
return Date.parse(withTimezone);
|
||||
}
|
||||
|
||||
function formatLocalTime(value: number | string | null | undefined): string {
|
||||
const timestamp = typeof value === 'number' ? value : parseServerTime(value);
|
||||
if (!Number.isFinite(timestamp)) return '-';
|
||||
return new Date(timestamp).toLocaleTimeString('zh-CN', {
|
||||
hour: '2-digit',
|
||||
minute: '2-digit',
|
||||
second: '2-digit',
|
||||
hour12: false,
|
||||
});
|
||||
}
|
||||
|
||||
function clampDimension(value: number) {
|
||||
if (!Number.isFinite(value)) return 1024;
|
||||
return Math.min(4096, Math.max(256, Math.round(value)));
|
||||
@@ -168,7 +227,62 @@ function templateSize(prompt: string) {
|
||||
return 'auto';
|
||||
}
|
||||
|
||||
function normalizeGalleryItems(items: GalleryItem[], limit = 48): PromptTemplate[] {
|
||||
function isPictureModel(model: string) {
|
||||
const normalized = model.trim().toLowerCase();
|
||||
if (!normalized) return false;
|
||||
const imageKeywords = [
|
||||
'gpt-image',
|
||||
'dall-e',
|
||||
'dalle',
|
||||
'midjourney',
|
||||
'recraft',
|
||||
'kolors',
|
||||
'cogview',
|
||||
'seedream',
|
||||
'sdxl',
|
||||
'stable-diffusion',
|
||||
'stablediffusion',
|
||||
'pixart',
|
||||
'playground',
|
||||
'firefly',
|
||||
'imagen',
|
||||
'flux',
|
||||
'image',
|
||||
'img',
|
||||
'canvas',
|
||||
'painting',
|
||||
'draw',
|
||||
];
|
||||
if (imageKeywords.some((keyword) => normalized.includes(keyword))) return true;
|
||||
const chatKeywords = [
|
||||
'claude',
|
||||
'chatgpt',
|
||||
'gpt-4',
|
||||
'gpt-3',
|
||||
'gpt-5',
|
||||
'o1',
|
||||
'o3',
|
||||
'o4',
|
||||
'gemini',
|
||||
'deepseek',
|
||||
'qwen',
|
||||
'llama',
|
||||
'mistral',
|
||||
'command',
|
||||
'kimi',
|
||||
'doubao',
|
||||
'glm',
|
||||
'hunyuan',
|
||||
'spark',
|
||||
'baichuan',
|
||||
'ernie',
|
||||
'wenxin',
|
||||
'grok',
|
||||
];
|
||||
return !chatKeywords.some((keyword) => normalized.includes(keyword));
|
||||
}
|
||||
|
||||
function normalizeGalleryItems(items: GalleryItem[], limit = TEMPLATE_LIMIT): PromptTemplate[] {
|
||||
return items
|
||||
.map<PromptTemplate | null>((item) => {
|
||||
const prompt = (item.prompt || '').trim();
|
||||
@@ -194,7 +308,7 @@ function normalizeGalleryItems(items: GalleryItem[], limit = 48): PromptTemplate
|
||||
.slice(0, limit);
|
||||
}
|
||||
|
||||
async function loadGalleryTemplatesDirectly(limit = 48): Promise<PromptTemplate[]> {
|
||||
async function loadGalleryTemplatesDirectly(limit = TEMPLATE_LIMIT): Promise<PromptTemplate[]> {
|
||||
const response = await fetch(TEMPLATE_GALLERY_JSON_URL, { cache: 'no-store' });
|
||||
if (!response.ok) throw new Error(`gallery HTTP ${response.status}`);
|
||||
const items = await response.json();
|
||||
@@ -221,20 +335,68 @@ export default function PictureAgent() {
|
||||
const [sourceFile, setSourceFile] = useState<File | null>(null);
|
||||
const [sourcePreview, setSourcePreview] = useState('');
|
||||
const [images, setImages] = useState<AgentImage[]>([]);
|
||||
const [jobs, setJobs] = useState<PictureJob[]>([]);
|
||||
const [selectedId, setSelectedId] = useState('');
|
||||
const [loading, setLoading] = useState(false);
|
||||
const [cancellingJobId, setCancellingJobId] = useState('');
|
||||
const [clearingJobs, setClearingJobs] = useState(false);
|
||||
const [jobDialogOpen, setJobDialogOpen] = useState(false);
|
||||
const [bootError, setBootError] = useState('');
|
||||
const [error, setError] = useState('');
|
||||
|
||||
const selectedImage = images.find((image) => image.id === selectedId) ?? images[0] ?? null;
|
||||
const activeJobs = jobs.filter((job) => job.status === 'pending' || job.status === 'running');
|
||||
const modeLabel = mode === 'edit' ? '图片编辑' : '文生图';
|
||||
const resultLabel = selectedImage?.mode === 'edit' ? '编辑结果' : '生成结果';
|
||||
const resolvedSize = buildSize(aspect, customWidth, customHeight);
|
||||
const latestJob = jobs[0] ?? null;
|
||||
const latestJobLabel = latestJob?.status === 'running'
|
||||
? '正在处理最新任务'
|
||||
: latestJob?.status === 'pending'
|
||||
? '最新任务已排队'
|
||||
: latestJob?.status === 'failed'
|
||||
? '最新任务失败'
|
||||
: latestJob?.status === 'cancelled'
|
||||
? '最新任务已取消'
|
||||
: latestJob?.status === 'succeeded'
|
||||
? '最新任务已完成'
|
||||
: '暂无最新任务';
|
||||
const terminalJobs = useMemo(
|
||||
() => jobs.filter((job) => job.status === 'succeeded' || job.status === 'failed' || job.status === 'cancelled'),
|
||||
[jobs],
|
||||
);
|
||||
|
||||
const authedFetch = useCallback((url: string, options: RequestInit & { timeoutMs?: number } = {}) => {
|
||||
return fetchSessionResponse(audience, url, options);
|
||||
}, [audience]);
|
||||
|
||||
const hydrateImagesFromJobs = useCallback((items: PictureJob[]) => {
|
||||
const completed = items.filter((job) => job.status === 'succeeded');
|
||||
if (!completed.length) return;
|
||||
setImages((current) => {
|
||||
const seen = new Set(current.map((image) => image.url));
|
||||
const next: AgentImage[] = [];
|
||||
for (const job of completed) {
|
||||
const urls = extractImageUrls(job.result || {});
|
||||
for (const url of urls) {
|
||||
if (seen.has(url)) continue;
|
||||
seen.add(url);
|
||||
next.push({
|
||||
id: makeId(),
|
||||
url,
|
||||
prompt: job.result?.prompt || '',
|
||||
model: job.result?.model || model || 'image',
|
||||
mode: job.result?.mode === 'edit' ? 'edit' : 'generate',
|
||||
createdAt: parseServerTime(job.finishedAt) || Date.now(),
|
||||
});
|
||||
}
|
||||
}
|
||||
if (!next.length) return current;
|
||||
setSelectedId((selected) => selected || next[0]?.id || '');
|
||||
return [...next, ...current].slice(0, 48);
|
||||
});
|
||||
}, [model]);
|
||||
|
||||
useEffect(() => {
|
||||
let ignore = false;
|
||||
async function load() {
|
||||
@@ -247,10 +409,11 @@ export default function PictureAgent() {
|
||||
setTemplatesLoading(true);
|
||||
setTemplatesError('');
|
||||
try {
|
||||
const [statusRes, modelsRes, templatesRes] = await Promise.all([
|
||||
const [statusRes, modelsRes, templatesRes, jobsRes] = await Promise.all([
|
||||
authedFetch('/api/picture-agent/status', { timeoutMs: 15_000 }),
|
||||
authedFetch('/api/picture-agent/models', { timeoutMs: 15_000 }),
|
||||
authedFetch('/api/picture-agent/templates?limit=48', { timeoutMs: 30_000 }),
|
||||
authedFetch(`/api/picture-agent/templates?limit=${TEMPLATE_LIMIT}`, { timeoutMs: 30_000 }),
|
||||
authedFetch('/api/picture-agent/jobs?limit=20', { timeoutMs: 15_000 }),
|
||||
]);
|
||||
if (!statusRes.ok) throw new Error(await parseError(statusRes));
|
||||
if (!modelsRes.ok) throw new Error(await parseError(modelsRes));
|
||||
@@ -262,17 +425,32 @@ export default function PictureAgent() {
|
||||
nextTemplates = Array.isArray(templatePayload?.templates) ? templatePayload.templates : [];
|
||||
} else {
|
||||
try {
|
||||
nextTemplates = await loadGalleryTemplatesDirectly(48);
|
||||
nextTemplates = await loadGalleryTemplatesDirectly(TEMPLATE_LIMIT);
|
||||
} catch {
|
||||
setTemplatesError(await parseError(templatesRes));
|
||||
}
|
||||
}
|
||||
let nextJobs: PictureJob[] = [];
|
||||
if (jobsRes.ok) {
|
||||
const jobsPayload = await jobsRes.json();
|
||||
nextJobs = Array.isArray(jobsPayload?.jobs) ? jobsPayload.jobs : [];
|
||||
}
|
||||
if (ignore) return;
|
||||
const availableModels = Array.isArray(nextModels?.models) ? nextModels.models : [];
|
||||
const availableModels = Array.isArray(nextModels?.models)
|
||||
? nextModels.models.filter((item: string) => isPictureModel(item))
|
||||
: [];
|
||||
setStatus(nextStatus);
|
||||
setModels(availableModels);
|
||||
setTemplates(nextTemplates);
|
||||
setModel((current) => current || nextModels?.defaultModel || availableModels[0] || 'gpt-image-2');
|
||||
setJobs(nextJobs);
|
||||
hydrateImagesFromJobs(nextJobs);
|
||||
setModel((current) => {
|
||||
if (current && isPictureModel(current)) return current;
|
||||
const fallback = nextModels?.defaultModel && isPictureModel(nextModels.defaultModel)
|
||||
? nextModels.defaultModel
|
||||
: availableModels[0];
|
||||
return fallback || 'gpt-image-2';
|
||||
});
|
||||
} catch (err: any) {
|
||||
if (!ignore) setBootError(err?.message || '生图助手初始化失败');
|
||||
} finally {
|
||||
@@ -283,7 +461,36 @@ export default function PictureAgent() {
|
||||
return () => {
|
||||
ignore = true;
|
||||
};
|
||||
}, [audience, authedFetch, token]);
|
||||
}, [audience, authedFetch, hydrateImagesFromJobs, token]);
|
||||
|
||||
useEffect(() => {
|
||||
if (!activeJobs.length || !token) return;
|
||||
let stopped = false;
|
||||
const timer = window.setInterval(async () => {
|
||||
try {
|
||||
const updated = await Promise.all(activeJobs.map(async (job) => {
|
||||
const response = await authedFetch(`/api/picture-agent/jobs/${encodeURIComponent(job.id)}`, { timeoutMs: 15_000 });
|
||||
if (!response.ok) return job;
|
||||
const payload = await response.json();
|
||||
return payload?.job || job;
|
||||
}));
|
||||
if (stopped) return;
|
||||
setJobs((current) => {
|
||||
const map = new Map(current.map((job) => [job.id, job]));
|
||||
for (const job of updated) map.set(job.id, job);
|
||||
const next = Array.from(map.values()).sort((a, b) => String(b.createdAt || '').localeCompare(String(a.createdAt || '')));
|
||||
hydrateImagesFromJobs(next);
|
||||
return next;
|
||||
});
|
||||
} catch {
|
||||
// Keep polling quietly; task status is also visible after refresh.
|
||||
}
|
||||
}, 1800);
|
||||
return () => {
|
||||
stopped = true;
|
||||
window.clearInterval(timer);
|
||||
};
|
||||
}, [activeJobs, authedFetch, hydrateImagesFromJobs, token]);
|
||||
|
||||
useEffect(() => {
|
||||
if (!sourceFile) {
|
||||
@@ -322,7 +529,7 @@ export default function PictureAgent() {
|
||||
timeoutMs: 180_000,
|
||||
});
|
||||
} else {
|
||||
response = await authedFetch('/api/picture-agent/images/generations', {
|
||||
response = await authedFetch('/api/picture-agent/jobs', {
|
||||
method: 'POST',
|
||||
body: JSON.stringify({
|
||||
model: model.trim(),
|
||||
@@ -330,12 +537,17 @@ export default function PictureAgent() {
|
||||
size: resolvedSize,
|
||||
quality,
|
||||
}),
|
||||
timeoutMs: 180_000,
|
||||
timeoutMs: 30_000,
|
||||
});
|
||||
}
|
||||
|
||||
if (!response.ok) throw new Error(await parseError(response));
|
||||
const payload = await response.json();
|
||||
if (mode === 'generate' && payload?.job) {
|
||||
setJobs((current) => [payload.job, ...current.filter((job) => job.id !== payload.job.id)].slice(0, 20));
|
||||
setError('');
|
||||
return;
|
||||
}
|
||||
const urls = extractImageUrls(payload);
|
||||
if (urls.length === 0) throw new Error('上游没有返回图片 URL 或 base64 图片。');
|
||||
|
||||
@@ -357,6 +569,56 @@ export default function PictureAgent() {
|
||||
}
|
||||
};
|
||||
|
||||
const refreshJobs = async () => {
|
||||
const response = await authedFetch('/api/picture-agent/jobs?limit=20', { timeoutMs: 15_000 });
|
||||
if (!response.ok) throw new Error(await parseError(response));
|
||||
const payload = await response.json();
|
||||
const nextJobs = Array.isArray(payload?.jobs) ? payload.jobs : [];
|
||||
setJobs(nextJobs);
|
||||
hydrateImagesFromJobs(nextJobs);
|
||||
};
|
||||
|
||||
const clearJobs = async () => {
|
||||
if (clearingJobs) return;
|
||||
setClearingJobs(true);
|
||||
setError('');
|
||||
try {
|
||||
const response = await authedFetch('/api/picture-agent/jobs/clear', {
|
||||
method: 'POST',
|
||||
timeoutMs: 15_000,
|
||||
});
|
||||
if (!response.ok) throw new Error(await parseError(response));
|
||||
await refreshJobs();
|
||||
} catch (err: any) {
|
||||
setError(err?.message || '清空日志失败');
|
||||
} finally {
|
||||
setClearingJobs(false);
|
||||
}
|
||||
};
|
||||
|
||||
const cancelJob = async (jobId: string) => {
|
||||
if (!jobId || cancellingJobId) return;
|
||||
setCancellingJobId(jobId);
|
||||
setError('');
|
||||
try {
|
||||
const response = await authedFetch(`/api/picture-agent/jobs/${encodeURIComponent(jobId)}/cancel`, {
|
||||
method: 'POST',
|
||||
timeoutMs: 15_000,
|
||||
});
|
||||
if (!response.ok) throw new Error(await parseError(response));
|
||||
const payload = await response.json();
|
||||
if (payload?.job) {
|
||||
setJobs((current) => current.map((job) => job.id === jobId ? payload.job : job));
|
||||
} else {
|
||||
await refreshJobs();
|
||||
}
|
||||
} catch (err: any) {
|
||||
setError(err?.message || '取消任务失败');
|
||||
} finally {
|
||||
setCancellingJobId('');
|
||||
}
|
||||
};
|
||||
|
||||
const handleDownload = (image: AgentImage | null) => {
|
||||
if (!image) return;
|
||||
const link = document.createElement('a');
|
||||
@@ -390,6 +652,59 @@ export default function PictureAgent() {
|
||||
setError('');
|
||||
};
|
||||
|
||||
const renderJobCard = (job: PictureJob) => {
|
||||
const total = Math.max(1, Number(job.progressTotal || 4));
|
||||
const progress = Math.min(total, Math.max(0, Number(job.progress || 0)));
|
||||
const percent = Math.round((progress / total) * 100);
|
||||
const hint = staleRunningHint(job);
|
||||
const isActive = job.status === 'pending' || job.status === 'running';
|
||||
const isFinished = !isActive;
|
||||
const statusText = job.status === 'succeeded'
|
||||
? '完成'
|
||||
: job.status === 'failed'
|
||||
? '失败'
|
||||
: job.status === 'cancelled'
|
||||
? '已取消'
|
||||
: `${percent}%`;
|
||||
return (
|
||||
<div key={job.id} className={`picture-agent-job picture-agent-job-${job.status}`}>
|
||||
<div className="picture-agent-job-row picture-agent-job-row-top">
|
||||
<div className="picture-agent-job-title-wrap">
|
||||
<strong>{job.result?.prompt || job.message || '生图任务'}</strong>
|
||||
<span>{job.model || job.result?.model || '未知模型'}</span>
|
||||
</div>
|
||||
<span className={`picture-agent-job-status picture-agent-job-status-${job.status}`}>{statusText}</span>
|
||||
</div>
|
||||
<div className="picture-agent-job-progress-line">
|
||||
<div className="picture-agent-job-progress-labels">
|
||||
<span>{job.message || job.id}</span>
|
||||
<em>{percent}%</em>
|
||||
</div>
|
||||
<div className="picture-agent-job-progress">
|
||||
<span style={{ width: `${percent}%` }} />
|
||||
</div>
|
||||
</div>
|
||||
<div className="picture-agent-job-meta">
|
||||
<span>{formatLocalTime(job.createdAt)}</span>
|
||||
{job.error ? <em>{job.error}</em> : <span>{isFinished ? '已结束' : '进行中'}</span>}
|
||||
</div>
|
||||
{hint ? <div className="picture-agent-job-hint">{hint}</div> : null}
|
||||
{isActive ? (
|
||||
<button
|
||||
type="button"
|
||||
className="picture-agent-job-cancel"
|
||||
disabled={cancellingJobId === job.id}
|
||||
onClick={() => void cancelJob(job.id)}
|
||||
>
|
||||
{cancellingJobId === job.id ? '取消中...' : '取消任务'}
|
||||
</button>
|
||||
) : (
|
||||
<div className="picture-agent-job-terminal">{job.result?.images?.length ? `结果 ${job.result.images.length} 张` : '可在历史结果中查看'}</div>
|
||||
)}
|
||||
</div>
|
||||
);
|
||||
};
|
||||
|
||||
if (bootError) {
|
||||
return (
|
||||
<div className="empty-state" style={{ minHeight: 320 }}>
|
||||
@@ -539,8 +854,18 @@ export default function PictureAgent() {
|
||||
|
||||
{error && <div className="alert alert-error">{error}</div>}
|
||||
<button type="button" className="btn btn-primary picture-agent-submit" disabled={!canSubmit} onClick={handleSubmit}>
|
||||
{loading ? '处理中...' : mode === 'edit' ? '编辑图片' : '生成图片'}
|
||||
{loading ? '提交中...' : mode === 'edit' ? '编辑图片' : '生成图片'}
|
||||
</button>
|
||||
|
||||
<div className="picture-agent-job-summary">
|
||||
<div>
|
||||
<strong>生图进度</strong>
|
||||
<span>{latestJobLabel} · {jobs.length} 条记录</span>
|
||||
</div>
|
||||
<button type="button" className="btn btn-ghost" onClick={() => setJobDialogOpen(true)}>
|
||||
查看进度
|
||||
</button>
|
||||
</div>
|
||||
</aside>
|
||||
|
||||
<section className="picture-agent-canvas">
|
||||
@@ -581,7 +906,7 @@ export default function PictureAgent() {
|
||||
<div className="picture-agent-board-meta">
|
||||
<div>
|
||||
<strong>{resultLabel}</strong>
|
||||
<span>{selectedImage.model} · {new Date(selectedImage.createdAt).toLocaleTimeString()}</span>
|
||||
<span>{selectedImage.model} · {formatLocalTime(selectedImage.createdAt)}</span>
|
||||
</div>
|
||||
<button type="button" className="btn btn-ghost" onClick={() => handleDownload(selectedImage)}>下载</button>
|
||||
</div>
|
||||
@@ -618,6 +943,58 @@ export default function PictureAgent() {
|
||||
</div>
|
||||
</section>
|
||||
</div>
|
||||
{jobDialogOpen && (
|
||||
<div
|
||||
className="picture-agent-job-dialog-backdrop"
|
||||
role="presentation"
|
||||
onMouseDown={(event) => {
|
||||
if (event.target === event.currentTarget) setJobDialogOpen(false);
|
||||
}}
|
||||
>
|
||||
<section
|
||||
className="picture-agent-job-dialog"
|
||||
role="dialog"
|
||||
aria-modal="true"
|
||||
aria-labelledby="picture-agent-job-dialog-title"
|
||||
>
|
||||
<div className="picture-agent-job-dialog-head">
|
||||
<div>
|
||||
<strong id="picture-agent-job-dialog-title">生图进度</strong>
|
||||
<span>{latestJobLabel} · {jobs.length} 条记录</span>
|
||||
</div>
|
||||
<div className="picture-agent-job-actions">
|
||||
<button
|
||||
type="button"
|
||||
className="btn btn-ghost"
|
||||
onClick={() => {
|
||||
refreshJobs().catch((err: any) => setError(err?.message || '刷新任务失败'));
|
||||
}}
|
||||
>
|
||||
刷新
|
||||
</button>
|
||||
<button
|
||||
type="button"
|
||||
className="btn btn-ghost"
|
||||
onClick={() => void clearJobs()}
|
||||
disabled={clearingJobs || terminalJobs.length === 0}
|
||||
>
|
||||
{clearingJobs ? '清理中...' : '清空日志'}
|
||||
</button>
|
||||
<button type="button" className="btn btn-ghost" onClick={() => setJobDialogOpen(false)}>
|
||||
关闭
|
||||
</button>
|
||||
</div>
|
||||
</div>
|
||||
{jobs.length === 0 ? (
|
||||
<div className="picture-agent-job-empty">提交后会先进入队列，完成后会自动出现在这里和使用记录里。</div>
|
||||
) : (
|
||||
<div className="picture-agent-job-list picture-agent-job-list-dialog">
|
||||
{jobs.map(renderJobCard)}
|
||||
</div>
|
||||
)}
|
||||
</section>
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
@@ -9,6 +9,7 @@ interface LogItem {
|
||||
status: string | null;
|
||||
http_status: number | null;
|
||||
latency_ms: number | null;
|
||||
first_byte_latency_ms?: number | null;
|
||||
prompt_tokens: number | null;
|
||||
completion_tokens: number | null;
|
||||
total_tokens: number | null;
|
||||
@@ -42,6 +43,23 @@ function latencyColor(ms: number | null): string {
|
||||
return 'var(--color-danger)';
|
||||
}
|
||||
|
||||
function displayLatency(log: LogItem): string {
|
||||
const firstByte = log.first_byte_latency_ms;
|
||||
const total = log.latency_ms;
|
||||
if (firstByte != null && Number.isFinite(firstByte)) {
|
||||
return `${firstByte}ms / ${total != null ? `${total}ms` : '-'}`;
|
||||
}
|
||||
return total != null ? `${total}ms` : '-';
|
||||
}
|
||||
|
||||
function latencyTextColor(log: LogItem): string {
|
||||
const firstByte = log.first_byte_latency_ms;
|
||||
if (firstByte != null && Number.isFinite(firstByte)) {
|
||||
return latencyColor(firstByte);
|
||||
}
|
||||
return latencyColor(log.latency_ms);
|
||||
}
|
||||
|
||||
export default function UsageRecords() {
|
||||
const [logs, setLogs] = useState<LogItem[]>([]);
|
||||
const [summary, setSummary] = useState<Summary | null>(null);
|
||||
@@ -181,6 +199,9 @@ export default function UsageRecords() {
|
||||
</button>
|
||||
)}
|
||||
</div>
|
||||
<div style={{ marginTop: 8, fontSize: 12, color: 'var(--color-text-muted)' }}>
|
||||
延迟优先显示首字延迟，Tokens 与费用依赖上游返回的 usage 数据。
|
||||
</div>
|
||||
</div>
|
||||
|
||||
{/* Table */}
|
||||
@@ -236,8 +257,8 @@ export default function UsageRecords() {
|
||||
<span style={{ marginLeft: 4, fontSize: 10, color: 'var(--color-text-muted)' }}>流式</span>
|
||||
)}
|
||||
</td>
|
||||
<td style={{ fontSize: 13, color: latencyColor(log.latency_ms), fontFamily: 'var(--font-mono)' }}>
|
||||
{log.latency_ms != null ? `${log.latency_ms}ms` : '-'}
|
||||
<td style={{ fontSize: 13, color: latencyTextColor(log), fontFamily: 'var(--font-mono)' }}>
|
||||
{displayLatency(log)}
|
||||
</td>
|
||||
<td style={{ fontSize: 12, fontFamily: 'var(--font-mono)' }}>
|
||||
{log.total_tokens != null ? log.total_tokens.toLocaleString() : '-'}
|
||||
|
||||
@@ -157,9 +157,9 @@ export default function UserDashboard({ onLogout }: { onLogout?: () => void }) {
|
||||
<section className="card" style={{ padding: '20px' }}>
|
||||
<h2 style={{ fontSize: '16px', fontWeight: 600, margin: '0 0 16px 0' }}>快速入口</h2>
|
||||
<div style={{ display: 'grid', gridTemplateColumns: '1fr 1fr', gap: '12px' }}>
|
||||
<Link to="/user/video-agent" className="btn btn-ghost" style={{ border: '1px solid var(--color-border)', height: 'auto', padding: '16px', display: 'flex', flexDirection: 'column', gap: '8px', alignItems: 'center' }}>
|
||||
<svg width="24" height="24" fill="none" viewBox="0 0 24 24" stroke="currentColor"><path strokeLinecap="round" strokeLinejoin="round" strokeWidth={1.5} d="M15 10l4.553-2.276A1 1 0 0121 8.618v6.764a1 1 0 01-1.447.894L15 14M5 18h8a2 2 0 002-2V8a2 2 0 00-2-2H5a2 2 0 00-2 2v8a2 2 0 002 2z" /></svg>
|
||||
<span>视频工作台</span>
|
||||
<Link to="/user/workflows" className="btn btn-ghost" style={{ border: '1px solid var(--color-border)', height: 'auto', padding: '16px', display: 'flex', flexDirection: 'column', gap: '8px', alignItems: 'center' }}>
|
||||
<svg width="24" height="24" fill="none" viewBox="0 0 24 24" stroke="currentColor"><path strokeLinecap="round" strokeLinejoin="round" strokeWidth={1.5} d="M4 6h5v5H4zM15 4h5v5h-5zM15 15h5v5h-5zM9 8h6M17.5 9v6M9 11l6 6" /></svg>
|
||||
<span>工作流</span>
|
||||
</Link>
|
||||
<Link to="/user/comfyui-agent" className="btn btn-ghost" style={{ border: '1px solid var(--color-border)', height: 'auto', padding: '16px', display: 'flex', flexDirection: 'column', gap: '8px', alignItems: 'center' }}>
|
||||
<svg width="24" height="24" fill="none" viewBox="0 0 24 24" stroke="currentColor"><path strokeLinecap="round" strokeLinejoin="round" strokeWidth={1.5} d="M14.752 11.168l-3.197-2.132A1 1 0 0010 9.87v4.263a1 1 0 001.555.832l3.197-2.132a1 1 0 000-1.664z" /><path strokeLinecap="round" strokeLinejoin="round" strokeWidth={1.5} d="M21 12a9 9 0 11-18 0 9 9 0 0118 0z" /></svg>
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
import React, { useState } from 'react';
|
||||
import { useNavigate, Link } from 'react-router-dom';
|
||||
import { commitSession } from '../authSession.js';
|
||||
import LoginShell from './auth/LoginShell.js';
|
||||
|
||||
export default function UserRegister({ onLogin }: { onLogin?: () => void }) {
|
||||
const navigate = useNavigate();
|
||||
@@ -46,26 +47,28 @@ export default function UserRegister({ onLogin }: { onLogin?: () => void }) {
|
||||
};
|
||||
|
||||
return (
|
||||
<div className="login-shell">
|
||||
<div className="login-surface animate-scale-in" style={{ maxWidth: 420 }}>
|
||||
<section className="login-brand-panel login-brand-panel-light" style={{ padding: '28px 24px' }}>
|
||||
<div className="login-brand-header" style={{ justifyContent: 'center' }}>
|
||||
<div className="brand-mark-frame brand-mark-frame-hero">
|
||||
<div className="brand-mark-canvas">
|
||||
<img src="/logo.png" alt="BoosAPI" className="login-brand-logo" />
|
||||
<LoginShell
|
||||
t={(text) => text}
|
||||
variant="user"
|
||||
eyebrow="用户注册"
|
||||
title="创建账号"
|
||||
description="注册后即可使用 API 代理服务。"
|
||||
error={error}
|
||||
footer={(
|
||||
<div style={{ display: 'flex', flexDirection: 'column', gap: 8, alignItems: 'center' }}>
|
||||
<span>
|
||||
已有账号？
|
||||
<Link to="/user/login" style={{ color: 'var(--color-primary)', textDecoration: 'none', marginLeft: 4 }}>
|
||||
登录
|
||||
</Link>
|
||||
</span>
|
||||
<span style={{ fontSize: 12, color: 'var(--color-text-muted)' }}>
|
||||
<Link to="/boos" style={{ color: 'var(--color-text-muted)' }}>← 管理员登录</Link>
|
||||
</span>
|
||||
</div>
|
||||
</div>
|
||||
<div className="login-brand-summary" style={{ textAlign: 'center' }}>
|
||||
<div className="login-brand-name">BoosAPI</div>
|
||||
<div className="login-brand-kicker">用户注册</div>
|
||||
</div>
|
||||
</div>
|
||||
</section>
|
||||
<section className="login-auth-stage">
|
||||
<div className="login-auth-panel">
|
||||
<h2 className="login-auth-title">创建账号</h2>
|
||||
<p className="login-auth-copy">注册后即可使用 API 代理服务。</p>
|
||||
<form onSubmit={handleSubmit}>
|
||||
)}
|
||||
>
|
||||
<form onSubmit={handleSubmit} className="login-register-form">
|
||||
<label className="login-auth-label" htmlFor="reg-username-input">用户名</label>
|
||||
<input
|
||||
id="reg-username-input"
|
||||
@@ -77,7 +80,7 @@ export default function UserRegister({ onLogin }: { onLogin?: () => void }) {
|
||||
autoComplete="username"
|
||||
required
|
||||
/>
|
||||
<label className="login-auth-label" htmlFor="reg-email-input" style={{ marginTop: 12 }}>邮箱</label>
|
||||
<label className="login-auth-label" htmlFor="reg-email-input">邮箱</label>
|
||||
<input
|
||||
id="reg-email-input"
|
||||
type="email"
|
||||
@@ -88,7 +91,7 @@ export default function UserRegister({ onLogin }: { onLogin?: () => void }) {
|
||||
autoComplete="email"
|
||||
required
|
||||
/>
|
||||
<label className="login-auth-label" htmlFor="reg-password-input" style={{ marginTop: 12 }}>密码</label>
|
||||
<label className="login-auth-label" htmlFor="reg-password-input">密码</label>
|
||||
<input
|
||||
id="reg-password-input"
|
||||
type="password"
|
||||
@@ -99,7 +102,7 @@ export default function UserRegister({ onLogin }: { onLogin?: () => void }) {
|
||||
autoComplete="new-password"
|
||||
required
|
||||
/>
|
||||
<label className="login-auth-label" htmlFor="reg-confirm-password-input" style={{ marginTop: 12 }}>确认密码</label>
|
||||
<label className="login-auth-label" htmlFor="reg-confirm-password-input">确认密码</label>
|
||||
<input
|
||||
id="reg-confirm-password-input"
|
||||
type="password"
|
||||
@@ -110,31 +113,24 @@ export default function UserRegister({ onLogin }: { onLogin?: () => void }) {
|
||||
autoComplete="new-password"
|
||||
required
|
||||
/>
|
||||
{error && (
|
||||
<div className="alert alert-error animate-shake" style={{ marginBottom: 12, marginTop: 8 }}>
|
||||
{error}
|
||||
</div>
|
||||
)}
|
||||
<button
|
||||
type="submit"
|
||||
disabled={loading || !username.trim() || !email.trim() || !password || !confirmPassword}
|
||||
className="btn btn-primary login-auth-submit"
|
||||
style={{ marginTop: 4 }}
|
||||
>
|
||||
{loading ? <><span className="spinner spinner-sm" style={{ borderTopColor: 'white', borderColor: 'rgba(255,255,255,0.3)' }} /> 注册中...</> : '注册'}
|
||||
{loading ? (
|
||||
<>
|
||||
<span
|
||||
className="spinner spinner-sm"
|
||||
style={{ borderTopColor: 'white', borderColor: 'rgba(255,255,255,0.3)' }}
|
||||
/>
|
||||
注册中...
|
||||
</>
|
||||
) : (
|
||||
'注册'
|
||||
)}
|
||||
</button>
|
||||
</form>
|
||||
<div className="login-auth-footer" style={{ flexDirection: 'column', gap: 8, marginTop: 16 }}>
|
||||
<span>
|
||||
已有账号？<Link to="/user/login" style={{ color: 'var(--color-primary)', textDecoration: 'none' }}>登录</Link>
|
||||
</span>
|
||||
<span style={{ fontSize: 12, color: 'var(--color-text-muted)' }}>
|
||||
<Link to="/" style={{ color: 'var(--color-text-muted)' }}>← 管理员登录</Link>
|
||||
</span>
|
||||
</div>
|
||||
</div>
|
||||
</section>
|
||||
</div>
|
||||
</div>
|
||||
</LoginShell>
|
||||
);
|
||||
}
|
||||
|
||||
@@ -0,0 +1,111 @@
|
||||
import React from 'react';
|
||||
|
||||
export default function Workflows() {
|
||||
const [ready, setReady] = React.useState(false);
|
||||
const [message, setMessage] = React.useState('正在加载工作流前端...');
|
||||
const iframeRef = React.useRef<HTMLIFrameElement | null>(null);
|
||||
|
||||
const openFullscreen = () => {
|
||||
window.open('/workflow-ui/', '_blank', 'noopener,noreferrer');
|
||||
};
|
||||
|
||||
const reloadFrame = () => {
|
||||
if (iframeRef.current) {
|
||||
iframeRef.current.src = '/workflow-ui/';
|
||||
}
|
||||
};
|
||||
|
||||
React.useEffect(() => {
|
||||
let cancelled = false;
|
||||
const controller = new AbortController();
|
||||
void fetch('/workflow-ui/', { signal: controller.signal })
|
||||
.then((response) => {
|
||||
if (cancelled) return;
|
||||
setReady(response.ok);
|
||||
setMessage(response.ok ? '' : 'Workflow Studio 前端未构建，请先执行 workflow-ui:install');
|
||||
})
|
||||
.catch((error) => {
|
||||
if (cancelled || error?.name === 'AbortError') return;
|
||||
setReady(false);
|
||||
setMessage('Workflow Studio 前端未构建，请先执行 workflow-ui:install');
|
||||
});
|
||||
return () => {
|
||||
cancelled = true;
|
||||
controller.abort();
|
||||
};
|
||||
}, []);
|
||||
|
||||
return (
|
||||
<div style={{ width: '100%', height: 'calc(100vh - 56px)', display: 'flex', flexDirection: 'column' }}>
|
||||
<div style={toolbarStyle}>
|
||||
<div style={toolbarTitleStyle}>工作流</div>
|
||||
<div style={toolbarActionsStyle}>
|
||||
<button type="button" onClick={reloadFrame} disabled={!ready} style={toolbarButtonStyle}>
|
||||
刷新
|
||||
</button>
|
||||
<button type="button" onClick={openFullscreen} style={primaryToolbarButtonStyle}>
|
||||
新窗口打开
|
||||
</button>
|
||||
</div>
|
||||
</div>
|
||||
{ready ? (
|
||||
<iframe
|
||||
ref={iframeRef}
|
||||
src="/workflow-ui/"
|
||||
title="Workflow Studio"
|
||||
allow="clipboard-read; clipboard-write"
|
||||
style={{ width: '100%', height: '100%', border: 'none', display: 'block' }}
|
||||
/>
|
||||
) : (
|
||||
<div style={{ flex: 1, display: 'grid', placeItems: 'center', padding: 24, textAlign: 'center' }}>
|
||||
<div>
|
||||
<div style={{ fontSize: 18, fontWeight: 700, marginBottom: 8 }}>Workflow Studio</div>
|
||||
<div style={{ color: 'var(--color-text-secondary)', lineHeight: 1.7 }}>{message}</div>
|
||||
</div>
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
const toolbarStyle: React.CSSProperties = {
|
||||
minHeight: 44,
|
||||
display: 'flex',
|
||||
alignItems: 'center',
|
||||
justifyContent: 'space-between',
|
||||
gap: 12,
|
||||
padding: '8px 12px',
|
||||
borderBottom: '1px solid var(--color-border-light)',
|
||||
background: 'var(--color-surface)',
|
||||
};
|
||||
|
||||
const toolbarTitleStyle: React.CSSProperties = {
|
||||
fontSize: 14,
|
||||
fontWeight: 800,
|
||||
color: 'var(--color-text-primary)',
|
||||
};
|
||||
|
||||
const toolbarActionsStyle: React.CSSProperties = {
|
||||
display: 'flex',
|
||||
alignItems: 'center',
|
||||
gap: 8,
|
||||
};
|
||||
|
||||
const toolbarButtonStyle: React.CSSProperties = {
|
||||
minHeight: 30,
|
||||
border: '1px solid var(--color-border)',
|
||||
borderRadius: 8,
|
||||
padding: '5px 10px',
|
||||
background: 'var(--color-bg)',
|
||||
color: 'var(--color-text-primary)',
|
||||
cursor: 'pointer',
|
||||
fontSize: 12,
|
||||
};
|
||||
|
||||
const primaryToolbarButtonStyle: React.CSSProperties = {
|
||||
...toolbarButtonStyle,
|
||||
borderColor: 'var(--color-primary)',
|
||||
background: 'var(--color-primary)',
|
||||
color: '#fff',
|
||||
fontWeight: 700,
|
||||
};
|
||||
@@ -2,13 +2,14 @@ import React, { CSSProperties, useCallback, useEffect, useMemo, useState } from
|
||||
import { useLocation } from 'react-router-dom';
|
||||
import { api, fetchSessionResponse, userApi, type WorkspaceAsset } from '../api.js';
|
||||
import type { SessionAudience } from '../authSession.js';
|
||||
import { formatDateTimeLocal } from './helpers/checkinLogTime.js';
|
||||
|
||||
const SOURCE_OPTIONS = [
|
||||
{ value: 'all', label: '全部来源' },
|
||||
{ value: 'picture-agent', label: '生图助手' },
|
||||
{ value: 'comfyui-agent', label: '视频助手' },
|
||||
{ value: 'video-agent', label: '视频工作台' },
|
||||
{ value: 'comfyui', label: 'ComfyUI' },
|
||||
{ value: 'workflow-studio', label: '工作流' },
|
||||
{ value: 'comfyui', label: '外部 ComfyUI' },
|
||||
];
|
||||
|
||||
const KIND_OPTIONS = [
|
||||
@@ -46,6 +47,10 @@ function assetTitle(asset: WorkspaceAsset) {
|
||||
return asset.title || asset.prompt || asset.fileName || asset.publicId;
|
||||
}
|
||||
|
||||
function assetCreatedAtLabel(asset: WorkspaceAsset) {
|
||||
return formatDateTimeLocal(asset.createdAt);
|
||||
}
|
||||
|
||||
export default function Workspace() {
|
||||
const location = useLocation();
|
||||
const audience: SessionAudience = location.pathname.startsWith('/user/') ? 'user' : 'admin';
|
||||
@@ -167,7 +172,7 @@ export default function Workspace() {
|
||||
<div>
|
||||
<h1 className="page-title">工作空间</h1>
|
||||
<p className="page-subtitle">
|
||||
{audience === 'user' ? '你的创作成品库' : '全部用户的创作成品库'} · 自动归档生图助手、视频助手和 ComfyUI 产出
|
||||
{audience === 'user' ? '你的创作成品库' : '全部用户的创作成品库'} · 自动归档生图助手、视频助手、工作流和外部 ComfyUI 产出
|
||||
</p>
|
||||
</div>
|
||||
<div className="page-actions">
|
||||
@@ -235,7 +240,7 @@ export default function Workspace() {
|
||||
{items.length === 0 ? (
|
||||
<div className="empty-state" style={{ minHeight: 360, gridColumn: '1 / -1' }}>
|
||||
<div className="empty-state-title">{loading ? '正在加载工作空间' : '暂无成品产出'}</div>
|
||||
<div className="empty-state-desc">生图助手、视频助手和 ComfyUI 的产出会归档到这里。</div>
|
||||
<div className="empty-state-desc">生图助手、视频助手、工作流和外部 ComfyUI 的产出会归档到这里。</div>
|
||||
</div>
|
||||
) : items.map((asset) => {
|
||||
const preview = previewUrls[asset.publicId] || asset.remoteUrl || '';
|
||||
@@ -268,7 +273,7 @@ export default function Workspace() {
|
||||
<span style={metaStyle}>{sourceLabel(asset.source)} · {asset.model || '未记录模型'}</span>
|
||||
<div style={cardFooterStyle}>
|
||||
<span style={statusPillStyle}>{statusLabel(asset.status)}</span>
|
||||
<span style={metaStyle}>{asset.createdAt || ''}</span>
|
||||
<span style={metaStyle}>{assetCreatedAtLabel(asset)}</span>
|
||||
</div>
|
||||
</div>
|
||||
</button>
|
||||
@@ -303,6 +308,7 @@ export default function Workspace() {
|
||||
<div><span>模型</span><strong>{selected.model || '未记录'}</strong></div>
|
||||
<div><span>大小</span><strong>{formatBytes(selected.sizeBytes)}</strong></div>
|
||||
<div><span>文件</span><strong>{selected.hasLocalFile ? '本地归档' : '远程引用'}</strong></div>
|
||||
<div><span>时间</span><strong>{assetCreatedAtLabel(selected)}</strong></div>
|
||||
</div>
|
||||
{selected.prompt && (
|
||||
<div style={promptPanelStyle}>
|
||||
|
||||
@@ -5,7 +5,7 @@ import { ToastProvider } from '../components/Toast.js';
|
||||
import Accounts from './Accounts.js';
|
||||
import { installAccountsSnapshotCompat } from './testApiCompat.js';
|
||||
|
||||
const { apiMock } = vi.hoisted(() => ({
|
||||
const { apiMock } = vi.hoisted(() => ({
|
||||
apiMock: {
|
||||
getAccounts: vi.fn(),
|
||||
getAccountsSnapshot: vi.fn(),
|
||||
@@ -13,8 +13,10 @@ const { apiMock } = vi.hoisted(() => ({
|
||||
getAccountTokens: vi.fn(),
|
||||
addAccount: vi.fn(),
|
||||
addAccountAvailableModels: vi.fn(),
|
||||
checkModels: vi.fn(),
|
||||
rebuildRoutes: vi.fn(),
|
||||
},
|
||||
}));
|
||||
}));
|
||||
|
||||
vi.mock('../api.js', () => ({
|
||||
api: apiMock,
|
||||
@@ -44,6 +46,15 @@ describe('Accounts CodingPlan initialization', () => {
|
||||
{ id: 10, name: 'Aliyun CodingPlan', url: 'https://coding.dashscope.aliyuncs.com/v1', platform: 'openai', status: 'active' },
|
||||
]);
|
||||
apiMock.getAccountTokens.mockResolvedValue([]);
|
||||
apiMock.checkModels.mockResolvedValue({
|
||||
message: '模型刷新成功',
|
||||
refresh: {
|
||||
status: 'success',
|
||||
modelCount: 3,
|
||||
modelsPreview: ['qwen3-coder-plus', 'qwen3.5-plus', 'glm-5'],
|
||||
},
|
||||
});
|
||||
apiMock.rebuildRoutes.mockResolvedValue({ success: true });
|
||||
apiMock.addAccount.mockResolvedValue({
|
||||
id: 18,
|
||||
siteId: 10,
|
||||
@@ -73,7 +84,7 @@ describe('Accounts CodingPlan initialization', () => {
|
||||
|
||||
const rendered = JSON.stringify(root.toJSON());
|
||||
expect(rendered).toContain('添加 API Key 连接');
|
||||
expect(rendered).toContain('阿里云 CodingPlan / OpenAI');
|
||||
expect(rendered).toContain('阿里云 Coding Plan (OpenAI)');
|
||||
expect(rendered).toContain('qwen3-coder-plus');
|
||||
|
||||
const checkedBoxes = root.root.findAll((node) => (
|
||||
@@ -110,8 +121,10 @@ describe('Accounts CodingPlan initialization', () => {
|
||||
}));
|
||||
expect(apiMock.addAccountAvailableModels).toHaveBeenCalledWith(
|
||||
18,
|
||||
expect.arrayContaining(['qwen3-coder-plus', 'qwen3.5-plus', 'glm-5']),
|
||||
['qwen3-coder-plus', 'qwen3.5-plus'],
|
||||
);
|
||||
expect(apiMock.checkModels).not.toHaveBeenCalled();
|
||||
expect(apiMock.rebuildRoutes).toHaveBeenCalledWith(false, false);
|
||||
} finally {
|
||||
root?.unmount();
|
||||
}
|
||||
@@ -127,6 +140,14 @@ describe('Accounts CodingPlan initialization', () => {
|
||||
tokenType: 'apikey',
|
||||
queued: false,
|
||||
});
|
||||
apiMock.checkModels.mockResolvedValue({
|
||||
message: '模型刷新成功',
|
||||
refresh: {
|
||||
status: 'success',
|
||||
modelCount: 2,
|
||||
modelsPreview: ['deepseek-chat', 'deepseek-reasoner'],
|
||||
},
|
||||
});
|
||||
|
||||
let root!: ReactTestRenderer;
|
||||
try {
|
||||
@@ -142,7 +163,7 @@ describe('Accounts CodingPlan initialization', () => {
|
||||
await flushMicrotasks();
|
||||
|
||||
const rendered = JSON.stringify(root.toJSON());
|
||||
expect(rendered).toContain('DeepSeek / OpenAI');
|
||||
expect(rendered).toContain('DeepSeek (OpenAI)');
|
||||
expect(rendered).toContain('deepseek-chat');
|
||||
expect(rendered).toContain('deepseek-reasoner');
|
||||
|
||||
@@ -175,6 +196,8 @@ describe('Accounts CodingPlan initialization', () => {
|
||||
19,
|
||||
['deepseek-chat', 'deepseek-reasoner'],
|
||||
);
|
||||
expect(apiMock.checkModels).not.toHaveBeenCalled();
|
||||
expect(apiMock.rebuildRoutes).toHaveBeenCalledWith(false, false);
|
||||
} finally {
|
||||
root?.unmount();
|
||||
}
|
||||
|
||||
@@ -31,14 +31,20 @@ async function readErrorReason(res: Response): Promise<string> {
|
||||
}
|
||||
}
|
||||
|
||||
type AdminLoginMode = 'token' | 'password';
|
||||
|
||||
export default function AdminLoginPage({ t, onLogin }: AdminLoginPageProps) {
|
||||
const [mode, setMode] = useState<AdminLoginMode>('token');
|
||||
const [token, setToken] = useState('');
|
||||
const [loading, setLoading] = useState(false);
|
||||
const [email, setEmail] = useState('');
|
||||
const [password, setPassword] = useState('');
|
||||
const [loadingMode, setLoadingMode] = useState<AdminLoginMode | null>(null);
|
||||
const [error, setError] = useState('');
|
||||
|
||||
const handleTokenLogin = async () => {
|
||||
const handleTokenLogin = async (event?: React.FormEvent) => {
|
||||
event?.preventDefault();
|
||||
if (!token.trim()) return;
|
||||
setLoading(true);
|
||||
setLoadingMode('token');
|
||||
setError('');
|
||||
try {
|
||||
const res = await fetch('/api/settings/auth/info', {
|
||||
@@ -54,22 +60,108 @@ export default function AdminLoginPage({ t, onLogin }: AdminLoginPageProps) {
|
||||
} catch {
|
||||
setError(t(resolveLoginErrorMessage('network', '')));
|
||||
} finally {
|
||||
setLoading(false);
|
||||
setLoadingMode(null);
|
||||
}
|
||||
};
|
||||
|
||||
const canSubmit = !!token.trim();
|
||||
const handlePasswordLogin = async (event?: React.FormEvent) => {
|
||||
event?.preventDefault();
|
||||
if (!email.trim() || !password) return;
|
||||
setLoadingMode('password');
|
||||
setError('');
|
||||
try {
|
||||
const res = await fetch('/api/auth/login', {
|
||||
method: 'POST',
|
||||
headers: { 'Content-Type': 'application/json' },
|
||||
body: JSON.stringify({ email: email.trim(), password }),
|
||||
});
|
||||
const text = await res.text();
|
||||
let parsed: {
|
||||
token?: unknown;
|
||||
user?: { role?: unknown };
|
||||
data?: { token?: unknown; user?: { role?: unknown } };
|
||||
message?: unknown;
|
||||
error?: unknown;
|
||||
} = {};
|
||||
if (text) {
|
||||
try {
|
||||
parsed = JSON.parse(text);
|
||||
} catch {
|
||||
parsed = { message: text };
|
||||
}
|
||||
}
|
||||
if (!res.ok) {
|
||||
const reason =
|
||||
typeof parsed.message === 'string'
|
||||
? parsed.message
|
||||
: typeof parsed.error === 'string'
|
||||
? parsed.error
|
||||
: text;
|
||||
setError(t(resolveLoginErrorMessage(res.status, reason)));
|
||||
return;
|
||||
}
|
||||
|
||||
const jwt = typeof parsed.token === 'string' ? parsed.token : parsed.data?.token;
|
||||
const role =
|
||||
typeof parsed.user?.role === 'string' ? parsed.user.role : parsed.data?.user?.role;
|
||||
if (typeof jwt !== 'string' || !jwt) {
|
||||
setError(t('登录响应缺少访问令牌。'));
|
||||
return;
|
||||
}
|
||||
if (role !== 'admin') {
|
||||
setError(t('该账号不是管理员账号。'));
|
||||
return;
|
||||
}
|
||||
commitSession('admin', jwt);
|
||||
onLogin(jwt);
|
||||
} catch {
|
||||
setError(t(resolveLoginErrorMessage('network', '')));
|
||||
} finally {
|
||||
setLoadingMode(null);
|
||||
}
|
||||
};
|
||||
|
||||
const isLoading = loadingMode !== null;
|
||||
const canTokenSubmit = !!token.trim();
|
||||
const canPasswordSubmit = !!email.trim() && !!password;
|
||||
|
||||
return (
|
||||
<LoginShell
|
||||
t={t}
|
||||
variant="admin"
|
||||
eyebrow={t('管理员入口')}
|
||||
title={t('登录')}
|
||||
description={t('请输入管理员令牌后继续。')}
|
||||
title={t('管理员登录')}
|
||||
description={t('使用管理员令牌，或管理员账号密码进入控制台。')}
|
||||
error={error}
|
||||
footer={<span>{t('管理员登录后继续。')}</span>}
|
||||
footer={<span>{t('管理员身份用于管理站点、模型路由、用户与计费。')}</span>}
|
||||
>
|
||||
<div className="login-mode-tabs" role="tablist" aria-label={t('管理员登录方式')}>
|
||||
<button
|
||||
type="button"
|
||||
className={`login-mode-tab${mode === 'token' ? ' active' : ''}`}
|
||||
onClick={() => {
|
||||
setMode('token');
|
||||
setError('');
|
||||
}}
|
||||
aria-selected={mode === 'token'}
|
||||
>
|
||||
Auth Token
|
||||
</button>
|
||||
<button
|
||||
type="button"
|
||||
className={`login-mode-tab${mode === 'password' ? ' active' : ''}`}
|
||||
onClick={() => {
|
||||
setMode('password');
|
||||
setError('');
|
||||
}}
|
||||
aria-selected={mode === 'password'}
|
||||
>
|
||||
{t('账号密码')}
|
||||
</button>
|
||||
</div>
|
||||
|
||||
{mode === 'token' ? (
|
||||
<form onSubmit={handleTokenLogin}>
|
||||
<label className="login-auth-label" htmlFor="admin-token-input">
|
||||
{t('管理员令牌')}
|
||||
</label>
|
||||
@@ -82,16 +174,15 @@ export default function AdminLoginPage({ t, onLogin }: AdminLoginPageProps) {
|
||||
setToken(e.target.value);
|
||||
setError('');
|
||||
}}
|
||||
onKeyDown={(e) => e.key === 'Enter' && handleTokenLogin()}
|
||||
className="login-auth-input"
|
||||
autoComplete="current-password"
|
||||
/>
|
||||
<button
|
||||
onClick={handleTokenLogin}
|
||||
disabled={loading || !canSubmit}
|
||||
type="submit"
|
||||
disabled={isLoading || !canTokenSubmit}
|
||||
className="btn btn-primary login-auth-submit"
|
||||
>
|
||||
{loading ? (
|
||||
{loadingMode === 'token' ? (
|
||||
<>
|
||||
<span
|
||||
className="spinner spinner-sm"
|
||||
@@ -100,12 +191,67 @@ export default function AdminLoginPage({ t, onLogin }: AdminLoginPageProps) {
|
||||
{t('验证中...')}
|
||||
</>
|
||||
) : (
|
||||
t('登录')
|
||||
t('进入控制台')
|
||||
)}
|
||||
</button>
|
||||
<div className="login-auth-note">
|
||||
{t('仅校验本地服务访问权限，不会把令牌发送到第三方。')}
|
||||
</div>
|
||||
</form>
|
||||
) : (
|
||||
<form onSubmit={handlePasswordLogin}>
|
||||
<label className="login-auth-label" htmlFor="admin-email-input">
|
||||
{t('管理员邮箱')}
|
||||
</label>
|
||||
<input
|
||||
id="admin-email-input"
|
||||
type="email"
|
||||
placeholder={t('管理员邮箱')}
|
||||
value={email}
|
||||
onChange={(e) => {
|
||||
setEmail(e.target.value);
|
||||
setError('');
|
||||
}}
|
||||
className="login-auth-input"
|
||||
autoComplete="username"
|
||||
/>
|
||||
<label className="login-auth-label" htmlFor="admin-password-input">
|
||||
{t('密码')}
|
||||
</label>
|
||||
<input
|
||||
id="admin-password-input"
|
||||
type="password"
|
||||
placeholder={t('密码')}
|
||||
value={password}
|
||||
onChange={(e) => {
|
||||
setPassword(e.target.value);
|
||||
setError('');
|
||||
}}
|
||||
className="login-auth-input"
|
||||
autoComplete="current-password"
|
||||
/>
|
||||
<button
|
||||
type="submit"
|
||||
disabled={isLoading || !canPasswordSubmit}
|
||||
className="btn btn-primary login-auth-submit"
|
||||
>
|
||||
{loadingMode === 'password' ? (
|
||||
<>
|
||||
<span
|
||||
className="spinner spinner-sm"
|
||||
style={{ borderTopColor: 'white', borderColor: 'rgba(255,255,255,0.3)' }}
|
||||
/>
|
||||
{t('验证中...')}
|
||||
</>
|
||||
) : (
|
||||
t('进入控制台')
|
||||
)}
|
||||
</button>
|
||||
<div className="login-auth-note">
|
||||
{t('只有 role=admin 的账号可以进入管理员控制台。')}
|
||||
</div>
|
||||
</form>
|
||||
)}
|
||||
</LoginShell>
|
||||
);
|
||||
}
|
||||
|
||||
@@ -1,5 +1,4 @@
|
||||
import React from 'react';
|
||||
import { SITE_DOCS_URL, SITE_GITHUB_URL } from '../../docsLink.js';
|
||||
|
||||
export type LoginShellVariant = 'admin' | 'user';
|
||||
|
||||
@@ -82,31 +81,6 @@ export default function LoginShell({
|
||||
</div>
|
||||
))}
|
||||
</div>
|
||||
<div className="login-brand-footer">
|
||||
<a
|
||||
href={SITE_GITHUB_URL}
|
||||
target="_blank"
|
||||
rel="noopener noreferrer"
|
||||
className="login-icon-link"
|
||||
aria-label="GitHub"
|
||||
title="GitHub"
|
||||
>
|
||||
<svg viewBox="0 0 24 24" aria-hidden="true" className="login-icon-link-svg">
|
||||
<path
|
||||
fill="currentColor"
|
||||
d="M12 2C6.48 2 2 6.59 2 12.25c0 4.53 2.87 8.38 6.84 9.73.5.1.68-.22.68-.49 0-.24-.01-1.04-.01-1.88-2.78.62-3.37-1.22-3.37-1.22-.45-1.2-1.11-1.52-1.11-1.52-.91-.64.07-.63.07-.63 1 .07 1.53 1.06 1.53 1.06.9 1.57 2.35 1.12 2.92.85.09-.67.35-1.12.64-1.38-2.22-.26-4.55-1.15-4.55-5.13 0-1.13.39-2.05 1.03-2.77-.1-.26-.45-1.31.1-2.73 0 0 .84-.28 2.75 1.06A9.3 9.3 0 0 1 12 6.91c.85 0 1.71.12 2.51.35 1.91-1.34 2.75-1.06 2.75-1.06.55 1.42.2 2.47.1 2.73.64.72 1.03 1.64 1.03 2.77 0 3.99-2.34 4.86-4.57 5.12.36.33.68.97.68 1.96 0 1.42-.01 2.56-.01 2.91 0 .27.18.59.69.49A10.27 10.27 0 0 0 22 12.25C22 6.59 17.52 2 12 2Z"
|
||||
/>
|
||||
</svg>
|
||||
</a>
|
||||
<a
|
||||
href={SITE_DOCS_URL}
|
||||
target="_blank"
|
||||
rel="noopener noreferrer"
|
||||
className="login-doc-link"
|
||||
>
|
||||
{t('部署文档')}
|
||||
</a>
|
||||
</div>
|
||||
</>
|
||||
)}
|
||||
</section>
|
||||
|
||||
@@ -0,0 +1,91 @@
|
||||
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
|
||||
import { act, create, type ReactTestInstance } from 'react-test-renderer';
|
||||
import { MemoryRouter } from 'react-router-dom';
|
||||
import { ToastProvider } from '../components/Toast.js';
|
||||
import Models from './Models.js';
|
||||
|
||||
const { apiMock, toastMock } = vi.hoisted(() => ({
|
||||
apiMock: {
|
||||
getModelsMarketplace: vi.fn(),
|
||||
},
|
||||
toastMock: {
|
||||
success: vi.fn(),
|
||||
error: vi.fn(),
|
||||
info: vi.fn(),
|
||||
},
|
||||
}));
|
||||
|
||||
vi.mock('../api.js', () => ({ api: apiMock }));
|
||||
vi.mock('../components/Toast.js', () => ({
|
||||
ToastProvider: ({ children }: { children: React.ReactNode }) => children,
|
||||
useToast: () => toastMock,
|
||||
}));
|
||||
|
||||
function collectText(node: ReactTestInstance): string {
|
||||
return (node.children || []).map((child) => (
|
||||
typeof child === 'string' ? child : collectText(child)
|
||||
)).join('');
|
||||
}
|
||||
|
||||
async function flushMicrotasks() {
|
||||
await act(async () => {
|
||||
await Promise.resolve();
|
||||
await Promise.resolve();
|
||||
});
|
||||
}
|
||||
|
||||
describe('Models copy action', () => {
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
apiMock.getModelsMarketplace.mockResolvedValue({
|
||||
models: [
|
||||
{
|
||||
name: 'doubao-seedance-2-0-260128',
|
||||
accountCount: 1,
|
||||
tokenCount: 1,
|
||||
avgLatency: null,
|
||||
successRate: null,
|
||||
description: null,
|
||||
tags: [],
|
||||
supportedEndpointTypes: [],
|
||||
pricingSources: [],
|
||||
accounts: [],
|
||||
},
|
||||
],
|
||||
});
|
||||
Object.defineProperty(globalThis.navigator, 'clipboard', {
|
||||
value: { writeText: vi.fn().mockResolvedValue(undefined) },
|
||||
configurable: true,
|
||||
});
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
vi.clearAllMocks();
|
||||
});
|
||||
|
||||
it('copies the model name from the action button', async () => {
|
||||
let root!: ReturnType<typeof create>;
|
||||
await act(async () => {
|
||||
root = create(
|
||||
<MemoryRouter initialEntries={['/models']}>
|
||||
<ToastProvider>
|
||||
<Models />
|
||||
</ToastProvider>
|
||||
</MemoryRouter>,
|
||||
);
|
||||
});
|
||||
await flushMicrotasks();
|
||||
|
||||
const copyBtn = root.root.find((node) => (
|
||||
node.type === 'button'
|
||||
&& typeof node.props.onClick === 'function'
|
||||
&& node.props['aria-label'] === '复制模型名'
|
||||
));
|
||||
await act(async () => {
|
||||
await copyBtn.props.onClick();
|
||||
});
|
||||
|
||||
expect(globalThis.navigator.clipboard.writeText).toHaveBeenCalledWith('doubao-seedance-2-0-260128');
|
||||
expect(toastMock.success).toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,92 @@
|
||||
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
|
||||
import { act, create, type ReactTestInstance } from 'react-test-renderer';
|
||||
import { MemoryRouter } from 'react-router-dom';
|
||||
import { ToastProvider } from '../components/Toast.js';
|
||||
import Models from './Models.js';
|
||||
|
||||
const { apiMock, userApiMock } = vi.hoisted(() => ({
|
||||
apiMock: {
|
||||
getModelsMarketplace: vi.fn(),
|
||||
},
|
||||
userApiMock: {
|
||||
getModelsMarketplace: vi.fn(),
|
||||
},
|
||||
}));
|
||||
|
||||
vi.mock('../api.js', () => ({
|
||||
api: apiMock,
|
||||
userApi: userApiMock,
|
||||
}));
|
||||
|
||||
function collectText(node: ReactTestInstance): string {
|
||||
return (node.children || []).map((child) => (
|
||||
typeof child === 'string' ? child : collectText(child)
|
||||
)).join('');
|
||||
}
|
||||
|
||||
async function flushMicrotasks() {
|
||||
await act(async () => {
|
||||
await Promise.resolve();
|
||||
await Promise.resolve();
|
||||
});
|
||||
}
|
||||
|
||||
describe('Models user view', () => {
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
userApiMock.getModelsMarketplace.mockResolvedValue({
|
||||
models: [
|
||||
{
|
||||
name: 'gpt-4o',
|
||||
accountCount: 1,
|
||||
tokenCount: 1,
|
||||
avgLatency: null,
|
||||
successRate: null,
|
||||
description: null,
|
||||
tags: [],
|
||||
supportedEndpointTypes: [],
|
||||
pricingSources: [],
|
||||
accounts: [
|
||||
{
|
||||
id: 1,
|
||||
site: 'Site A',
|
||||
username: 'tester',
|
||||
latency: null,
|
||||
balance: 1,
|
||||
tokens: [{ id: 1, name: 'default', isDefault: true }],
|
||||
},
|
||||
],
|
||||
},
|
||||
],
|
||||
});
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
vi.clearAllMocks();
|
||||
});
|
||||
|
||||
it('hides vendor and account details for user path', async () => {
|
||||
let root!: ReturnType<typeof create>;
|
||||
await act(async () => {
|
||||
root = create(
|
||||
<MemoryRouter initialEntries={['/user/models']}>
|
||||
<ToastProvider>
|
||||
<Models />
|
||||
</ToastProvider>
|
||||
</MemoryRouter>,
|
||||
);
|
||||
});
|
||||
await flushMicrotasks();
|
||||
|
||||
const text = collectText(root.root);
|
||||
expect(text).toContain('模型广场');
|
||||
expect(text).not.toContain('供应商');
|
||||
expect(text).not.toContain('排序方式');
|
||||
expect(text).not.toContain('账号数');
|
||||
expect(text).not.toContain('令牌数');
|
||||
expect(text).not.toContain('Site A');
|
||||
expect(text).not.toContain('tester');
|
||||
expect(userApiMock.getModelsMarketplace).toHaveBeenCalled();
|
||||
expect(apiMock.getModelsMarketplace).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,925 @@
|
||||
use crate::{
|
||||
error::{AppError, AppResult},
|
||||
services::workspace,
|
||||
workflow::{
|
||||
graph,
|
||||
models::{NodeExecution, NodeOutputStore, RunContext, WorkflowGraph, WorkflowNode},
|
||||
storage,
|
||||
},
|
||||
AppState,
|
||||
};
|
||||
use serde_json::{json, Map, Value};
|
||||
use std::{collections::HashMap, sync::Arc};
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum WorkflowProgressEvent {
|
||||
Started {
|
||||
prompt_id: String,
|
||||
total: usize,
|
||||
},
|
||||
NodeStarted {
|
||||
prompt_id: String,
|
||||
node_id: String,
|
||||
index: usize,
|
||||
total: usize,
|
||||
},
|
||||
NodeCompleted {
|
||||
prompt_id: String,
|
||||
node_id: String,
|
||||
index: usize,
|
||||
total: usize,
|
||||
output: Value,
|
||||
},
|
||||
Finished {
|
||||
prompt_id: String,
|
||||
success: bool,
|
||||
error: Option<String>,
|
||||
},
|
||||
}
|
||||
|
||||
pub type WorkflowProgressCallback = Arc<dyn Fn(WorkflowProgressEvent) + Send + Sync>;
|
||||
|
||||
pub async fn run_workflow(
|
||||
state: &AppState,
|
||||
workflow_id: i64,
|
||||
graph_value: Value,
|
||||
input_overrides: Value,
|
||||
ctx: RunContext,
|
||||
) -> AppResult<String> {
|
||||
run_workflow_with_progress(state, workflow_id, graph_value, input_overrides, ctx, None).await
|
||||
}
|
||||
|
||||
pub async fn run_workflow_with_progress(
|
||||
state: &AppState,
|
||||
workflow_id: i64,
|
||||
graph_value: Value,
|
||||
input_overrides: Value,
|
||||
ctx: RunContext,
|
||||
progress: Option<WorkflowProgressCallback>,
|
||||
) -> AppResult<String> {
|
||||
let graph: WorkflowGraph = serde_json::from_value(graph_value.clone())
|
||||
.map_err(|err| AppError::BadRequest(format!("invalid workflow graph: {err}")))?;
|
||||
let order = graph::validate_and_sort(&graph)?;
|
||||
let total = order.len().max(1);
|
||||
let (run_db_id, run_public_id) =
|
||||
storage::insert_run(state, workflow_id, ctx.user_id, &input_overrides).await?;
|
||||
emit_progress(
|
||||
&progress,
|
||||
WorkflowProgressEvent::Started {
|
||||
prompt_id: run_public_id.clone(),
|
||||
total,
|
||||
},
|
||||
);
|
||||
let nodes = graph
|
||||
.nodes
|
||||
.iter()
|
||||
.map(|node| (node.id.clone(), node.clone()))
|
||||
.collect::<HashMap<_, _>>();
|
||||
let mut store = NodeOutputStore::new();
|
||||
|
||||
for (index, node_id) in order.into_iter().enumerate() {
|
||||
let Some(node) = nodes.get(&node_id) else {
|
||||
continue;
|
||||
};
|
||||
emit_progress(
|
||||
&progress,
|
||||
WorkflowProgressEvent::NodeStarted {
|
||||
prompt_id: run_public_id.clone(),
|
||||
node_id: node.id.clone(),
|
||||
index,
|
||||
total,
|
||||
},
|
||||
);
|
||||
let inputs = collect_inputs(&graph, &store, node, &input_overrides)?;
|
||||
let node_run_id =
|
||||
storage::insert_node_run_started(state, run_db_id, &node.id, &node.node_type, &inputs)
|
||||
.await?;
|
||||
match execute_node(state, &ctx, node, inputs.clone()).await {
|
||||
Ok(result) => {
|
||||
storage::complete_node_run(state, node_run_id, &result.outputs).await?;
|
||||
emit_progress(
|
||||
&progress,
|
||||
WorkflowProgressEvent::NodeCompleted {
|
||||
prompt_id: run_public_id.clone(),
|
||||
node_id: node.id.clone(),
|
||||
index: index + 1,
|
||||
total,
|
||||
output: result.outputs.clone(),
|
||||
},
|
||||
);
|
||||
store.insert(node.id.clone(), result.outputs);
|
||||
}
|
||||
Err(err) => {
|
||||
let message = err.to_string();
|
||||
let _ = storage::fail_node_run(state, node_run_id, &message).await;
|
||||
let _ = storage::fail_run(state, run_db_id, &message).await;
|
||||
emit_progress(
|
||||
&progress,
|
||||
WorkflowProgressEvent::Finished {
|
||||
prompt_id: run_public_id.clone(),
|
||||
success: false,
|
||||
error: Some(message),
|
||||
},
|
||||
);
|
||||
return Err(err);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let final_outputs = json!({ "nodes": store });
|
||||
storage::complete_run(state, run_db_id, &final_outputs).await?;
|
||||
emit_progress(
|
||||
&progress,
|
||||
WorkflowProgressEvent::Finished {
|
||||
prompt_id: run_public_id.clone(),
|
||||
success: true,
|
||||
error: None,
|
||||
},
|
||||
);
|
||||
Ok(run_public_id)
|
||||
}
|
||||
|
||||
fn emit_progress(progress: &Option<WorkflowProgressCallback>, event: WorkflowProgressEvent) {
|
||||
if let Some(callback) = progress {
|
||||
callback(event);
|
||||
}
|
||||
}
|
||||
|
||||
fn collect_inputs(
|
||||
graph: &WorkflowGraph,
|
||||
store: &NodeOutputStore,
|
||||
node: &WorkflowNode,
|
||||
input_overrides: &Value,
|
||||
) -> AppResult<Value> {
|
||||
let mut inputs = Map::new();
|
||||
for edge in &graph.edges {
|
||||
if graph::endpoint_node(&edge.to)? != node.id {
|
||||
continue;
|
||||
}
|
||||
let from_node = graph::endpoint_node(&edge.from)?;
|
||||
let from_port = graph::endpoint_port(&edge.from)?;
|
||||
let to_port = graph::endpoint_port(&edge.to)?;
|
||||
if let Some(value) = store
|
||||
.get(from_node)
|
||||
.and_then(|outputs| outputs.get(from_port))
|
||||
{
|
||||
inputs.insert(to_port.to_string(), value.clone());
|
||||
}
|
||||
}
|
||||
if let Some(overrides) = input_overrides.as_object() {
|
||||
for (key, value) in overrides {
|
||||
if let Some((override_node, port)) = key.split_once('.') {
|
||||
if override_node == node.id {
|
||||
inputs.insert(port.to_string(), value.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(Value::Object(inputs))
|
||||
}
|
||||
|
||||
async fn execute_node(
|
||||
state: &AppState,
|
||||
ctx: &RunContext,
|
||||
node: &WorkflowNode,
|
||||
inputs: Value,
|
||||
) -> AppResult<NodeExecution> {
|
||||
match node.node_type.as_str() {
|
||||
"input.text" => {
|
||||
let text = inputs
|
||||
.get("value")
|
||||
.or_else(|| node.params.get("value"))
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default()
|
||||
.to_string();
|
||||
Ok(NodeExecution {
|
||||
inputs,
|
||||
outputs: json!({ "text": text }),
|
||||
})
|
||||
}
|
||||
"boosapi.chat_completion" => chat_completion(state, ctx, node, inputs, false).await,
|
||||
"boosapi.comfyui_workflow" => comfyui_workflow(ctx, node, inputs).await,
|
||||
"comfyui.submit_workflow" => submit_comfyui_workflow(state, node, inputs).await,
|
||||
"boosapi.prompt_rewrite" => chat_completion(state, ctx, node, inputs, true).await,
|
||||
"boosapi.text_to_image" => text_to_image(state, ctx, node, inputs).await,
|
||||
"boosapi.text_to_video" => text_to_video(ctx, node, inputs).await,
|
||||
"boosapi.image_to_video" => image_to_video(ctx, node, inputs).await,
|
||||
"boosapi.video_status" => video_status(ctx, node, inputs).await,
|
||||
"CheckpointLoaderSimple" => {
|
||||
let model = param_str(node, "ckpt_name", "gpt-image-2");
|
||||
Ok(NodeExecution {
|
||||
inputs,
|
||||
outputs: json!({
|
||||
"MODEL": { "model": model },
|
||||
"CLIP": { "model": model },
|
||||
"VAE": { "model": model }
|
||||
}),
|
||||
})
|
||||
}
|
||||
"CLIPTextEncode" => {
|
||||
let text = param_str(node, "text", "");
|
||||
let model = inputs
|
||||
.pointer("/clip/model")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or("gpt-image-2")
|
||||
.to_string();
|
||||
Ok(NodeExecution {
|
||||
inputs,
|
||||
outputs: json!({ "CONDITIONING": { "text": text, "model": model } }),
|
||||
})
|
||||
}
|
||||
"EmptyLatentImage" => Ok(NodeExecution {
|
||||
inputs,
|
||||
outputs: json!({
|
||||
"LATENT": {
|
||||
"width": param_i64(node, "width", 1024),
|
||||
"height": param_i64(node, "height", 1024),
|
||||
"batch_size": param_i64(node, "batch_size", 1)
|
||||
}
|
||||
}),
|
||||
}),
|
||||
"KSampler" => {
|
||||
let model = inputs
|
||||
.pointer("/model/model")
|
||||
.and_then(Value::as_str)
|
||||
.or_else(|| inputs.pointer("/positive/model").and_then(Value::as_str))
|
||||
.unwrap_or("gpt-image-2")
|
||||
.to_string();
|
||||
let prompt = inputs
|
||||
.pointer("/positive/text")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default()
|
||||
.to_string();
|
||||
let negative_prompt = inputs
|
||||
.pointer("/negative/text")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default()
|
||||
.to_string();
|
||||
let width = inputs
|
||||
.pointer("/latent_image/width")
|
||||
.and_then(Value::as_i64)
|
||||
.unwrap_or(1024);
|
||||
let height = inputs
|
||||
.pointer("/latent_image/height")
|
||||
.and_then(Value::as_i64)
|
||||
.unwrap_or(1024);
|
||||
Ok(NodeExecution {
|
||||
inputs,
|
||||
outputs: json!({
|
||||
"LATENT": {
|
||||
"model": model,
|
||||
"prompt": prompt,
|
||||
"negative_prompt": negative_prompt,
|
||||
"width": width,
|
||||
"height": height,
|
||||
"steps": param_i64(node, "steps", 20),
|
||||
"cfg": param_number(node, "cfg", 8.0),
|
||||
"seed": param_i64(node, "seed", 0)
|
||||
}
|
||||
}),
|
||||
})
|
||||
}
|
||||
"VAEDecode" => vae_decode_to_image(state, ctx, inputs).await,
|
||||
"SaveImage" | "PreviewImage" => {
|
||||
let image = inputs
|
||||
.get("images")
|
||||
.cloned()
|
||||
.or_else(|| inputs.get("image").cloned())
|
||||
.or_else(|| inputs.get("IMAGE").cloned())
|
||||
.ok_or_else(|| {
|
||||
AppError::BadRequest(format!("{} missing images input", node.id))
|
||||
})?;
|
||||
Ok(NodeExecution {
|
||||
inputs,
|
||||
outputs: json!({ "IMAGE": image, "image": image }),
|
||||
})
|
||||
}
|
||||
"output.preview_image" => {
|
||||
let image = inputs
|
||||
.get("image")
|
||||
.cloned()
|
||||
.or_else(|| inputs.get("images").cloned())
|
||||
.or_else(|| inputs.get("IMAGE").cloned())
|
||||
.ok_or_else(|| AppError::BadRequest(format!("{} missing image input", node.id)))?;
|
||||
Ok(NodeExecution {
|
||||
inputs,
|
||||
outputs: json!({ "image": image }),
|
||||
})
|
||||
}
|
||||
"output.preview_video" => {
|
||||
let video = inputs
|
||||
.get("video")
|
||||
.cloned()
|
||||
.or_else(|| inputs.get("videos").cloned())
|
||||
.or_else(|| inputs.get("VIDEO").cloned())
|
||||
.ok_or_else(|| AppError::BadRequest(format!("{} missing video input", node.id)))?;
|
||||
Ok(NodeExecution {
|
||||
inputs,
|
||||
outputs: json!({ "video": video }),
|
||||
})
|
||||
}
|
||||
"output.save_file" => {
|
||||
let file = inputs
|
||||
.get("file")
|
||||
.cloned()
|
||||
.or_else(|| inputs.get("image").cloned())
|
||||
.or_else(|| inputs.get("images").and_then(|images| images.get(0)).cloned())
|
||||
.ok_or_else(|| AppError::BadRequest(format!("{} missing file input", node.id)))?;
|
||||
Ok(NodeExecution {
|
||||
inputs,
|
||||
outputs: json!({ "file": file }),
|
||||
})
|
||||
}
|
||||
"util.text_join" => {
|
||||
let prefix = param_str(node, "prefix", "");
|
||||
let suffix = param_str(node, "suffix", "");
|
||||
let separator = param_str(node, "separator", "\n");
|
||||
let text = value_to_string(inputs.get("text"));
|
||||
let joined = [prefix, text, suffix]
|
||||
.into_iter()
|
||||
.filter(|value| !value.is_empty())
|
||||
.collect::<Vec<_>>()
|
||||
.join(&separator);
|
||||
Ok(NodeExecution {
|
||||
inputs,
|
||||
outputs: json!({ "text": joined }),
|
||||
})
|
||||
}
|
||||
"util.json_extract" => {
|
||||
let mut current = inputs.get("json").cloned().unwrap_or(Value::Null);
|
||||
for segment in param_str(node, "path", "")
|
||||
.split('.')
|
||||
.filter(|v| !v.is_empty())
|
||||
{
|
||||
current = if let Ok(index) = segment.parse::<usize>() {
|
||||
current.get(index).cloned().unwrap_or(Value::Null)
|
||||
} else {
|
||||
current.get(segment).cloned().unwrap_or(Value::Null)
|
||||
};
|
||||
}
|
||||
Ok(NodeExecution {
|
||||
inputs,
|
||||
outputs: json!({ "value": current, "text": value_to_string(Some(¤t)) }),
|
||||
})
|
||||
}
|
||||
_ => Err(AppError::BadRequest(format!(
|
||||
"unsupported node type: {}",
|
||||
node.node_type
|
||||
))),
|
||||
}
|
||||
}
|
||||
|
||||
async fn comfyui_workflow(
|
||||
ctx: &RunContext,
|
||||
node: &WorkflowNode,
|
||||
inputs: Value,
|
||||
) -> AppResult<NodeExecution> {
|
||||
let text = value_to_string(inputs.get("text"));
|
||||
if text.trim().is_empty() {
|
||||
return Err(AppError::BadRequest(format!(
|
||||
"{} missing text input",
|
||||
node.id
|
||||
)));
|
||||
}
|
||||
let system = param_str(
|
||||
node,
|
||||
"system",
|
||||
"Generate a valid ComfyUI workflow JSON object. Return JSON only.",
|
||||
);
|
||||
let body = json!({
|
||||
"model": param_str(node, "model", "gpt-5.4"),
|
||||
"messages": [
|
||||
{ "role": "system", "content": system },
|
||||
{ "role": "user", "content": text }
|
||||
],
|
||||
"stream": false,
|
||||
"temperature": param_number(node, "temperature", 0.3)
|
||||
});
|
||||
let raw = post_json(ctx, "/chat/completions", &body).await?;
|
||||
let content = raw
|
||||
.pointer("/choices/0/message/content")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default()
|
||||
.to_string();
|
||||
let workflow = parse_jsonish(&content).unwrap_or_else(|| {
|
||||
json!({
|
||||
"boosapi": {
|
||||
"kind": "comfyui-workflow-draft",
|
||||
"prompt": text,
|
||||
"raw": content
|
||||
},
|
||||
"nodes": {}
|
||||
})
|
||||
});
|
||||
Ok(NodeExecution {
|
||||
inputs,
|
||||
outputs: json!({ "workflow": workflow, "text": content, "raw": raw }),
|
||||
})
|
||||
}
|
||||
|
||||
async fn submit_comfyui_workflow(
|
||||
state: &AppState,
|
||||
node: &WorkflowNode,
|
||||
inputs: Value,
|
||||
) -> AppResult<NodeExecution> {
|
||||
let workflow = inputs
|
||||
.get("workflow")
|
||||
.cloned()
|
||||
.ok_or_else(|| AppError::BadRequest(format!("{} missing workflow input", node.id)))?;
|
||||
if param_bool(node, "dry_run", true) {
|
||||
return Ok(NodeExecution {
|
||||
inputs,
|
||||
outputs: json!({
|
||||
"prompt_id": null,
|
||||
"response": { "dryRun": true },
|
||||
"workflow": workflow,
|
||||
}),
|
||||
});
|
||||
}
|
||||
let body = json!({
|
||||
"client_id": param_str(node, "client_id", "boosapi-workflow-studio"),
|
||||
"prompt": workflow,
|
||||
});
|
||||
let response = reqwest::Client::new()
|
||||
.post(format!("{}/api/comfyui/prompt", ctx_server_base(state)))
|
||||
.json(&body)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|err| AppError::UpstreamError(err.to_string()))?;
|
||||
let status = response.status();
|
||||
let text = response
|
||||
.text()
|
||||
.await
|
||||
.map_err(|err| AppError::UpstreamError(err.to_string()))?;
|
||||
let value = serde_json::from_str::<Value>(&text).unwrap_or_else(|_| json!({ "body": text }));
|
||||
if !status.is_success() {
|
||||
return Err(AppError::UpstreamError(value.to_string()));
|
||||
}
|
||||
Ok(NodeExecution {
|
||||
inputs,
|
||||
outputs: json!({
|
||||
"prompt_id": value.get("prompt_id").and_then(Value::as_str),
|
||||
"response": value,
|
||||
"workflow": body.get("prompt").cloned().unwrap_or(Value::Null),
|
||||
}),
|
||||
})
|
||||
}
|
||||
|
||||
async fn chat_completion(
|
||||
_state: &AppState,
|
||||
ctx: &RunContext,
|
||||
node: &WorkflowNode,
|
||||
inputs: Value,
|
||||
rewrite: bool,
|
||||
) -> AppResult<NodeExecution> {
|
||||
let text = value_to_string(inputs.get("text"));
|
||||
if text.trim().is_empty() {
|
||||
return Err(AppError::BadRequest(format!(
|
||||
"{} missing text input",
|
||||
node.id
|
||||
)));
|
||||
}
|
||||
let model = param_str(node, "model", "gpt-5.4");
|
||||
let system = if rewrite {
|
||||
format!(
|
||||
"Rewrite the user input as a concise, high quality prompt. Style requirement: {}. Return only the rewritten prompt.",
|
||||
param_str(node, "style", "cinematic English image prompt")
|
||||
)
|
||||
} else {
|
||||
param_str(node, "system", "")
|
||||
};
|
||||
let mut messages = Vec::new();
|
||||
if !system.trim().is_empty() {
|
||||
messages.push(json!({ "role": "system", "content": system }));
|
||||
}
|
||||
messages.push(json!({ "role": "user", "content": text }));
|
||||
let body = json!({
|
||||
"model": model,
|
||||
"messages": messages,
|
||||
"stream": false,
|
||||
"temperature": param_number(node, "temperature", 0.7)
|
||||
});
|
||||
let raw = post_json(ctx, "/chat/completions", &body).await?;
|
||||
let content = raw
|
||||
.pointer("/choices/0/message/content")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default()
|
||||
.to_string();
|
||||
let outputs = if rewrite {
|
||||
json!({ "prompt": content, "raw": raw })
|
||||
} else {
|
||||
json!({ "text": content, "raw": raw })
|
||||
};
|
||||
Ok(NodeExecution { inputs, outputs })
|
||||
}
|
||||
|
||||
async fn text_to_image(
|
||||
state: &AppState,
|
||||
ctx: &RunContext,
|
||||
node: &WorkflowNode,
|
||||
inputs: Value,
|
||||
) -> AppResult<NodeExecution> {
|
||||
let prompt = value_to_string(inputs.get("prompt"));
|
||||
if prompt.trim().is_empty() {
|
||||
return Err(AppError::BadRequest(format!(
|
||||
"{} missing prompt input",
|
||||
node.id
|
||||
)));
|
||||
}
|
||||
let width = param_i64(node, "width", 1024);
|
||||
let height = param_i64(node, "height", 1024);
|
||||
let size = param_str(node, "size", "");
|
||||
let body = json!({
|
||||
"model": param_str(node, "model", "gpt-image-2"),
|
||||
"prompt": prompt,
|
||||
"n": 1,
|
||||
"size": if size.trim().is_empty() { format!("{width}x{height}") } else { size },
|
||||
});
|
||||
execute_image_generation(state, ctx, inputs, body).await
|
||||
}
|
||||
|
||||
async fn text_to_video(
|
||||
ctx: &RunContext,
|
||||
node: &WorkflowNode,
|
||||
inputs: Value,
|
||||
) -> AppResult<NodeExecution> {
|
||||
let prompt = value_to_string(inputs.get("prompt"));
|
||||
if prompt.trim().is_empty() {
|
||||
return Err(AppError::BadRequest(format!(
|
||||
"{} missing prompt input",
|
||||
node.id
|
||||
)));
|
||||
}
|
||||
let body = video_request_body(node, prompt, None);
|
||||
execute_video_generation(ctx, inputs, body).await
|
||||
}
|
||||
|
||||
async fn image_to_video(
|
||||
ctx: &RunContext,
|
||||
node: &WorkflowNode,
|
||||
inputs: Value,
|
||||
) -> AppResult<NodeExecution> {
|
||||
let prompt = value_to_string(inputs.get("prompt"));
|
||||
if prompt.trim().is_empty() {
|
||||
return Err(AppError::BadRequest(format!(
|
||||
"{} missing prompt input",
|
||||
node.id
|
||||
)));
|
||||
}
|
||||
let image_url = param_str(node, "image_url", "");
|
||||
let image_url = if image_url.trim().is_empty() {
|
||||
extract_media_url(inputs.get("image"))
|
||||
} else {
|
||||
Some(image_url)
|
||||
};
|
||||
let body = video_request_body(node, prompt, image_url);
|
||||
execute_video_generation(ctx, inputs, body).await
|
||||
}
|
||||
|
||||
async fn video_status(
|
||||
ctx: &RunContext,
|
||||
node: &WorkflowNode,
|
||||
inputs: Value,
|
||||
) -> AppResult<NodeExecution> {
|
||||
let task_id = value_to_string(inputs.get("task_id"));
|
||||
if task_id.trim().is_empty() {
|
||||
return Err(AppError::BadRequest(format!(
|
||||
"{} missing task_id input",
|
||||
node.id
|
||||
)));
|
||||
}
|
||||
let model = param_str(node, "model", "seedance-1-0-lite-t2v");
|
||||
let path = format!(
|
||||
"/videos/{}?model={}",
|
||||
simple_urlencode(task_id.trim()),
|
||||
simple_urlencode(model.trim())
|
||||
);
|
||||
let raw = get_json(ctx, &path).await?;
|
||||
let video = video_output_from_response(&raw);
|
||||
Ok(NodeExecution {
|
||||
inputs,
|
||||
outputs: json!({ "video": video, "raw": raw }),
|
||||
})
|
||||
}
|
||||
|
||||
fn video_request_body(node: &WorkflowNode, prompt: String, image_url: Option<String>) -> Value {
|
||||
let mut body = Map::new();
|
||||
body.insert("model".to_string(), json!(param_str(node, "model", "seedance-1-0-lite-t2v")));
|
||||
body.insert("prompt".to_string(), json!(prompt));
|
||||
for key in ["ratio", "resolution"] {
|
||||
let value = param_str(node, key, "");
|
||||
if !value.trim().is_empty() {
|
||||
body.insert(key.to_string(), json!(value));
|
||||
}
|
||||
}
|
||||
for key in ["duration", "fps"] {
|
||||
if let Some(value) = node.params.get(key).and_then(Value::as_i64) {
|
||||
body.insert(key.to_string(), json!(value));
|
||||
} else if let Some(value) = node.params.get(key).and_then(Value::as_f64) {
|
||||
body.insert(key.to_string(), json!(value));
|
||||
}
|
||||
}
|
||||
if let Some(url) = image_url.filter(|value| !value.trim().is_empty()) {
|
||||
body.insert("image_url".to_string(), json!(url));
|
||||
}
|
||||
Value::Object(body)
|
||||
}
|
||||
|
||||
async fn execute_video_generation(
|
||||
ctx: &RunContext,
|
||||
inputs: Value,
|
||||
body: Value,
|
||||
) -> AppResult<NodeExecution> {
|
||||
let raw = post_json(ctx, "/videos", &body).await?;
|
||||
let task_id = extract_task_id(&raw);
|
||||
let video = video_output_from_response(&raw);
|
||||
Ok(NodeExecution {
|
||||
inputs,
|
||||
outputs: json!({
|
||||
"task_id": task_id,
|
||||
"video": video,
|
||||
"raw": raw
|
||||
}),
|
||||
})
|
||||
}
|
||||
|
||||
async fn vae_decode_to_image(
|
||||
state: &AppState,
|
||||
ctx: &RunContext,
|
||||
inputs: Value,
|
||||
) -> AppResult<NodeExecution> {
|
||||
let samples = inputs.get("samples").cloned().unwrap_or(Value::Null);
|
||||
let vae = inputs.get("vae").cloned().unwrap_or(Value::Null);
|
||||
let prompt = samples
|
||||
.get("prompt")
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or_default()
|
||||
.to_string();
|
||||
if prompt.trim().is_empty() {
|
||||
return Err(AppError::BadRequest(
|
||||
"VAEDecode missing positive prompt from KSampler".into(),
|
||||
));
|
||||
}
|
||||
let model = samples
|
||||
.get("model")
|
||||
.and_then(Value::as_str)
|
||||
.or_else(|| vae.get("model").and_then(Value::as_str))
|
||||
.unwrap_or("gpt-image-2");
|
||||
let width = samples
|
||||
.get("width")
|
||||
.and_then(Value::as_i64)
|
||||
.unwrap_or(1024);
|
||||
let height = samples
|
||||
.get("height")
|
||||
.and_then(Value::as_i64)
|
||||
.unwrap_or(1024);
|
||||
let mut prompt_text = prompt;
|
||||
if let Some(negative) = samples.get("negative_prompt").and_then(Value::as_str) {
|
||||
if !negative.trim().is_empty() {
|
||||
prompt_text.push_str("\nNegative prompt: ");
|
||||
prompt_text.push_str(negative.trim());
|
||||
}
|
||||
}
|
||||
let body = json!({
|
||||
"model": model,
|
||||
"prompt": prompt_text,
|
||||
"n": 1,
|
||||
"size": format!("{width}x{height}"),
|
||||
});
|
||||
execute_image_generation(state, ctx, inputs, body).await
|
||||
}
|
||||
|
||||
async fn execute_image_generation(
|
||||
state: &AppState,
|
||||
ctx: &RunContext,
|
||||
inputs: Value,
|
||||
body: Value,
|
||||
) -> AppResult<NodeExecution> {
|
||||
let raw = post_json(ctx, "/images/generations", &body).await?;
|
||||
let assets = workspace::record_image_assets_from_response(
|
||||
state,
|
||||
ctx.user_id,
|
||||
"workflow-studio",
|
||||
body.get("prompt")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::to_string),
|
||||
body.get("model")
|
||||
.and_then(Value::as_str)
|
||||
.map(str::to_string),
|
||||
&raw,
|
||||
)
|
||||
.await;
|
||||
let images = raw.get("data").cloned().unwrap_or_else(|| json!([]));
|
||||
let first = images.get(0).cloned().unwrap_or(Value::Null);
|
||||
let comfy_images = assets
|
||||
.iter()
|
||||
.map(|asset| {
|
||||
json!({
|
||||
"filename": asset.file_name.as_deref().unwrap_or(&asset.public_id),
|
||||
"subfolder": "",
|
||||
"type": "output",
|
||||
})
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
Ok(NodeExecution {
|
||||
inputs,
|
||||
outputs: json!({
|
||||
"image": first,
|
||||
"images": images,
|
||||
"comfy_images": comfy_images,
|
||||
"raw": raw
|
||||
}),
|
||||
})
|
||||
}
|
||||
|
||||
async fn post_json(ctx: &RunContext, path: &str, body: &Value) -> AppResult<Value> {
|
||||
let url = format!("{}{}", ctx.base_url.trim_end_matches('/'), path);
|
||||
let response = reqwest::Client::new()
|
||||
.post(url)
|
||||
.bearer_auth(&ctx.bearer)
|
||||
.json(body)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|err| AppError::UpstreamError(err.to_string()))?;
|
||||
let status = response.status();
|
||||
let text = response
|
||||
.text()
|
||||
.await
|
||||
.map_err(|err| AppError::UpstreamError(err.to_string()))?;
|
||||
let json = serde_json::from_str::<Value>(&text).map_err(|err| {
|
||||
AppError::UpstreamError(format!(
|
||||
"invalid JSON from BoosAPI workflow node: {err}; body: {}",
|
||||
text.chars().take(500).collect::<String>()
|
||||
))
|
||||
})?;
|
||||
if !status.is_success() {
|
||||
return Err(AppError::UpstreamError(json.to_string()));
|
||||
}
|
||||
Ok(json)
|
||||
}
|
||||
|
||||
async fn get_json(ctx: &RunContext, path: &str) -> AppResult<Value> {
|
||||
let url = format!("{}{}", ctx.base_url.trim_end_matches('/'), path);
|
||||
let response = reqwest::Client::new()
|
||||
.get(url)
|
||||
.bearer_auth(&ctx.bearer)
|
||||
.send()
|
||||
.await
|
||||
.map_err(|err| AppError::UpstreamError(err.to_string()))?;
|
||||
let status = response.status();
|
||||
let text = response
|
||||
.text()
|
||||
.await
|
||||
.map_err(|err| AppError::UpstreamError(err.to_string()))?;
|
||||
let json = serde_json::from_str::<Value>(&text).map_err(|err| {
|
||||
AppError::UpstreamError(format!(
|
||||
"invalid JSON from BoosAPI workflow node: {err}; body: {}",
|
||||
text.chars().take(500).collect::<String>()
|
||||
))
|
||||
})?;
|
||||
if !status.is_success() {
|
||||
return Err(AppError::UpstreamError(json.to_string()));
|
||||
}
|
||||
Ok(json)
|
||||
}
|
||||
|
||||
fn extract_task_id(value: &Value) -> Value {
|
||||
for pointer in ["/id", "/task_id", "/data/id", "/data/task_id", "/result/id", "/result/task_id"] {
|
||||
if let Some(id) = value.pointer(pointer).and_then(Value::as_str) {
|
||||
if !id.trim().is_empty() {
|
||||
return json!(id);
|
||||
}
|
||||
}
|
||||
}
|
||||
Value::Null
|
||||
}
|
||||
|
||||
fn video_output_from_response(value: &Value) -> Value {
|
||||
if let Some(url) = find_url_by_extension(value, &[".mp4", ".webm", ".mov"])
|
||||
.or_else(|| find_url_key(value, &["video_url", "videoUrl", "url"]))
|
||||
{
|
||||
json!({ "url": url })
|
||||
} else {
|
||||
Value::Null
|
||||
}
|
||||
}
|
||||
|
||||
fn extract_media_url(value: Option<&Value>) -> Option<String> {
|
||||
let value = value?;
|
||||
find_url_key(value, &["url", "remote_url", "remoteUrl", "image_url", "imageUrl"])
|
||||
}
|
||||
|
||||
fn find_url_key(value: &Value, keys: &[&str]) -> Option<String> {
|
||||
match value {
|
||||
Value::Object(map) => {
|
||||
for key in keys {
|
||||
if let Some(url) = map
|
||||
.get(*key)
|
||||
.and_then(Value::as_str)
|
||||
.filter(|url| url.starts_with("http://") || url.starts_with("https://"))
|
||||
{
|
||||
return Some(url.to_string());
|
||||
}
|
||||
}
|
||||
map.values().find_map(|child| find_url_key(child, keys))
|
||||
}
|
||||
Value::Array(items) => items.iter().find_map(|child| find_url_key(child, keys)),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn find_url_by_extension(value: &Value, extensions: &[&str]) -> Option<String> {
|
||||
match value {
|
||||
Value::String(text)
|
||||
if (text.starts_with("http://") || text.starts_with("https://"))
|
||||
&& extensions.iter().any(|ext| text.to_ascii_lowercase().contains(ext)) =>
|
||||
{
|
||||
Some(text.clone())
|
||||
}
|
||||
Value::Object(map) => map
|
||||
.values()
|
||||
.find_map(|child| find_url_by_extension(child, extensions)),
|
||||
Value::Array(items) => items
|
||||
.iter()
|
||||
.find_map(|child| find_url_by_extension(child, extensions)),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn simple_urlencode(value: &str) -> String {
|
||||
value
|
||||
.replace('%', "%25")
|
||||
.replace(' ', "%20")
|
||||
.replace('?', "%3F")
|
||||
.replace('&', "%26")
|
||||
.replace('#', "%23")
|
||||
.replace('/', "%2F")
|
||||
}
|
||||
|
||||
fn param_str(node: &WorkflowNode, key: &str, default: &str) -> String {
|
||||
node.params
|
||||
.get(key)
|
||||
.and_then(Value::as_str)
|
||||
.unwrap_or(default)
|
||||
.to_string()
|
||||
}
|
||||
|
||||
fn param_i64(node: &WorkflowNode, key: &str, default: i64) -> i64 {
|
||||
node.params
|
||||
.get(key)
|
||||
.and_then(Value::as_i64)
|
||||
.unwrap_or(default)
|
||||
}
|
||||
|
||||
fn param_number(node: &WorkflowNode, key: &str, default: f64) -> f64 {
|
||||
node.params
|
||||
.get(key)
|
||||
.and_then(Value::as_f64)
|
||||
.unwrap_or(default)
|
||||
}
|
||||
|
||||
fn param_bool(node: &WorkflowNode, key: &str, default: bool) -> bool {
|
||||
node.params
|
||||
.get(key)
|
||||
.and_then(Value::as_bool)
|
||||
.unwrap_or(default)
|
||||
}
|
||||
|
||||
fn parse_jsonish(text: &str) -> Option<Value> {
|
||||
let trimmed = text.trim();
|
||||
if trimmed.is_empty() {
|
||||
return None;
|
||||
}
|
||||
let stripped = trimmed
|
||||
.strip_prefix("```json")
|
||||
.or_else(|| trimmed.strip_prefix("```"))
|
||||
.unwrap_or(trimmed)
|
||||
.trim()
|
||||
.strip_suffix("```")
|
||||
.unwrap_or_else(|| {
|
||||
trimmed
|
||||
.strip_prefix("```json")
|
||||
.or_else(|| trimmed.strip_prefix("```"))
|
||||
.unwrap_or(trimmed)
|
||||
.trim()
|
||||
})
|
||||
.trim();
|
||||
serde_json::from_str(stripped).ok()
|
||||
}
|
||||
|
||||
fn ctx_server_base(state: &AppState) -> String {
|
||||
let host = if state.config.server.host == "0.0.0.0" || state.config.server.host == "::" {
|
||||
"127.0.0.1".to_string()
|
||||
} else {
|
||||
state.config.server.host.clone()
|
||||
};
|
||||
format!("http://{}:{}", host, state.config.server.port)
|
||||
}
|
||||
|
||||
fn value_to_string(value: Option<&Value>) -> String {
|
||||
match value {
|
||||
Some(Value::String(text)) => text.clone(),
|
||||
Some(Value::Null) | None => String::new(),
|
||||
Some(other) => other.to_string(),
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,94 @@
|
||||
use crate::{
|
||||
error::{AppError, AppResult},
|
||||
workflow::{models::WorkflowGraph, registry},
|
||||
};
|
||||
use std::collections::{HashMap, VecDeque};
|
||||
|
||||
pub fn validate_and_sort(graph: &WorkflowGraph) -> AppResult<Vec<String>> {
|
||||
if graph.nodes.is_empty() {
|
||||
return Err(AppError::BadRequest(
|
||||
"workflow must contain at least one node".into(),
|
||||
));
|
||||
}
|
||||
let mut indegree: HashMap<String, usize> = HashMap::new();
|
||||
let mut outgoing: HashMap<String, Vec<String>> = HashMap::new();
|
||||
for node in &graph.nodes {
|
||||
let node_id = node.id.trim();
|
||||
if node_id.is_empty() {
|
||||
return Err(AppError::BadRequest("node id is required".into()));
|
||||
}
|
||||
if indegree.insert(node.id.clone(), 0).is_some() {
|
||||
return Err(AppError::BadRequest(format!(
|
||||
"duplicate node id: {}",
|
||||
node.id
|
||||
)));
|
||||
}
|
||||
if !registry::is_known_node(&node.node_type) {
|
||||
return Err(AppError::BadRequest(format!(
|
||||
"unknown node type: {}",
|
||||
node.node_type
|
||||
)));
|
||||
}
|
||||
}
|
||||
for edge in &graph.edges {
|
||||
let from = endpoint_node(&edge.from)?;
|
||||
let to = endpoint_node(&edge.to)?;
|
||||
if !indegree.contains_key(from) {
|
||||
return Err(AppError::BadRequest(format!(
|
||||
"edge references missing node: {from}"
|
||||
)));
|
||||
}
|
||||
if !indegree.contains_key(to) {
|
||||
return Err(AppError::BadRequest(format!(
|
||||
"edge references missing node: {to}"
|
||||
)));
|
||||
}
|
||||
outgoing
|
||||
.entry(from.to_string())
|
||||
.or_default()
|
||||
.push(to.to_string());
|
||||
*indegree.entry(to.to_string()).or_default() += 1;
|
||||
}
|
||||
|
||||
let mut queue = graph
|
||||
.nodes
|
||||
.iter()
|
||||
.filter(|node| indegree.get(&node.id).copied().unwrap_or(0) == 0)
|
||||
.map(|node| node.id.clone())
|
||||
.collect::<VecDeque<_>>();
|
||||
let mut order = Vec::with_capacity(graph.nodes.len());
|
||||
while let Some(node_id) = queue.pop_front() {
|
||||
order.push(node_id.clone());
|
||||
if let Some(next) = outgoing.get(&node_id) {
|
||||
for target in next {
|
||||
let entry = indegree.entry(target.clone()).or_default();
|
||||
*entry = entry.saturating_sub(1);
|
||||
if *entry == 0 {
|
||||
queue.push_back(target.clone());
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
if order.len() != graph.nodes.len() {
|
||||
return Err(AppError::BadRequest(
|
||||
"workflow graph contains a cycle".into(),
|
||||
));
|
||||
}
|
||||
Ok(order)
|
||||
}
|
||||
|
||||
pub fn endpoint_node(endpoint: &str) -> AppResult<&str> {
|
||||
endpoint
|
||||
.split_once('.')
|
||||
.map(|(node, _)| node)
|
||||
.filter(|value| !value.trim().is_empty())
|
||||
.ok_or_else(|| AppError::BadRequest(format!("invalid endpoint: {endpoint}")))
|
||||
}
|
||||
|
||||
pub fn endpoint_port(endpoint: &str) -> AppResult<&str> {
|
||||
endpoint
|
||||
.split_once('.')
|
||||
.map(|(_, port)| port)
|
||||
.filter(|value| !value.trim().is_empty())
|
||||
.ok_or_else(|| AppError::BadRequest(format!("invalid endpoint: {endpoint}")))
|
||||
}
|
||||
@@ -0,0 +1,6 @@
|
||||
pub mod executor;
|
||||
pub mod graph;
|
||||
pub mod models;
|
||||
pub mod registry;
|
||||
pub mod routes;
|
||||
pub mod storage;
|
||||
@@ -0,0 +1,155 @@
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::{json, Map, Value};
|
||||
use std::collections::HashMap;
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct WorkflowGraph {
|
||||
#[serde(default = "default_version")]
|
||||
pub version: String,
|
||||
#[serde(default)]
|
||||
pub name: Option<String>,
|
||||
#[serde(default)]
|
||||
pub nodes: Vec<WorkflowNode>,
|
||||
#[serde(default)]
|
||||
pub edges: Vec<WorkflowEdge>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct WorkflowNode {
|
||||
pub id: String,
|
||||
#[serde(rename = "type")]
|
||||
pub node_type: String,
|
||||
#[serde(default)]
|
||||
pub position: Value,
|
||||
#[serde(default)]
|
||||
pub params: Map<String, Value>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct WorkflowEdge {
|
||||
pub from: String,
|
||||
pub to: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct WorkflowSummary {
|
||||
pub id: i64,
|
||||
pub public_id: String,
|
||||
pub user_id: Option<i64>,
|
||||
pub name: String,
|
||||
pub description: Option<String>,
|
||||
pub graph: Value,
|
||||
pub is_template: bool,
|
||||
pub visibility: String,
|
||||
pub created_at: String,
|
||||
pub updated_at: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct WorkflowRunSummary {
|
||||
pub id: i64,
|
||||
pub public_id: String,
|
||||
pub workflow_id: i64,
|
||||
pub workflow_public_id: Option<String>,
|
||||
pub user_id: Option<i64>,
|
||||
pub status: String,
|
||||
pub input_overrides: Value,
|
||||
pub final_outputs: Option<Value>,
|
||||
pub error_message: Option<String>,
|
||||
pub started_at: Option<String>,
|
||||
pub completed_at: Option<String>,
|
||||
pub created_at: String,
|
||||
pub node_runs: Vec<WorkflowNodeRunSummary>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct WorkflowNodeRunSummary {
|
||||
pub id: i64,
|
||||
pub node_id: String,
|
||||
pub node_type: String,
|
||||
pub status: String,
|
||||
pub inputs: Value,
|
||||
pub outputs: Option<Value>,
|
||||
pub error_message: Option<String>,
|
||||
pub started_at: Option<String>,
|
||||
pub completed_at: Option<String>,
|
||||
pub created_at: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub struct NodeTypeDefinition {
|
||||
#[serde(rename = "type")]
|
||||
pub node_type: &'static str,
|
||||
pub title: &'static str,
|
||||
pub category: &'static str,
|
||||
pub description: &'static str,
|
||||
pub inputs: Value,
|
||||
pub params: Value,
|
||||
pub outputs: Value,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct RunContext {
|
||||
pub user_id: Option<i64>,
|
||||
pub bearer: String,
|
||||
pub base_url: String,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct NodeExecution {
|
||||
pub inputs: Value,
|
||||
pub outputs: Value,
|
||||
}
|
||||
|
||||
pub type NodeOutputStore = HashMap<String, Value>;
|
||||
|
||||
pub fn default_graph(name: &str) -> Value {
|
||||
json!({
|
||||
"version": default_version(),
|
||||
"name": name,
|
||||
"nodes": [
|
||||
{
|
||||
"id": "input_1",
|
||||
"type": "input.text",
|
||||
"position": { "x": 80, "y": 120 },
|
||||
"params": { "value": "生成一个电影感雨夜城市短镜头，包含一张 1024x1024 的起始画面" }
|
||||
},
|
||||
{
|
||||
"id": "workflow_1",
|
||||
"type": "boosapi.comfyui_workflow",
|
||||
"position": { "x": 360, "y": 120 },
|
||||
"params": {
|
||||
"model": "gpt-5.4",
|
||||
"system": "Generate a valid ComfyUI workflow JSON object. Return JSON only."
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": "submit_1",
|
||||
"type": "comfyui.submit_workflow",
|
||||
"position": { "x": 660, "y": 120 },
|
||||
"params": {
|
||||
"client_id": "boosapi-workflow-studio",
|
||||
"dry_run": true
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": "save_1",
|
||||
"type": "output.save_file",
|
||||
"position": { "x": 960, "y": 120 },
|
||||
"params": {}
|
||||
}
|
||||
],
|
||||
"edges": [
|
||||
{ "from": "input_1.text", "to": "workflow_1.text" },
|
||||
{ "from": "workflow_1.workflow", "to": "submit_1.workflow" },
|
||||
{ "from": "submit_1.workflow", "to": "save_1.file" }
|
||||
]
|
||||
})
|
||||
}
|
||||
|
||||
fn default_version() -> String {
|
||||
"1.0".to_string()
|
||||
}
|
||||
@@ -0,0 +1,291 @@
|
||||
use crate::workflow::models::NodeTypeDefinition;
|
||||
use serde_json::json;
|
||||
|
||||
fn core_node_types() -> Vec<NodeTypeDefinition> {
|
||||
vec![
|
||||
NodeTypeDefinition {
|
||||
node_type: "CheckpointLoaderSimple",
|
||||
title: "CheckpointLoaderSimple",
|
||||
category: "loaders",
|
||||
description: "Loads a checkpoint model and provides model, clip and vae outputs.",
|
||||
inputs: json!({}),
|
||||
params: json!({
|
||||
"ckpt_name": { "ui": "combo", "type": "COMBO", "default": "gpt-image-2" }
|
||||
}),
|
||||
outputs: json!({
|
||||
"MODEL": { "type": "MODEL" },
|
||||
"CLIP": { "type": "CLIP" },
|
||||
"VAE": { "type": "VAE" }
|
||||
}),
|
||||
},
|
||||
NodeTypeDefinition {
|
||||
node_type: "CLIPTextEncode",
|
||||
title: "CLIPTextEncode",
|
||||
category: "conditioning",
|
||||
description: "Encodes text into conditioning.",
|
||||
inputs: json!({
|
||||
"clip": { "type": "CLIP", "required": false }
|
||||
}),
|
||||
params: json!({
|
||||
"text": { "ui": "textarea", "type": "STRING", "default": "" }
|
||||
}),
|
||||
outputs: json!({ "CONDITIONING": { "type": "CONDITIONING" } }),
|
||||
},
|
||||
NodeTypeDefinition {
|
||||
node_type: "EmptyLatentImage",
|
||||
title: "EmptyLatentImage",
|
||||
category: "latent",
|
||||
description: "Creates an empty latent image.",
|
||||
inputs: json!({}),
|
||||
params: json!({
|
||||
"width": { "ui": "number", "type": "INT", "default": 512, "min": 64, "max": 4096 },
|
||||
"height": { "ui": "number", "type": "INT", "default": 512, "min": 64, "max": 4096 },
|
||||
"batch_size": { "ui": "number", "type": "INT", "default": 1, "min": 1, "max": 64 }
|
||||
}),
|
||||
outputs: json!({ "LATENT": { "type": "LATENT" } }),
|
||||
},
|
||||
NodeTypeDefinition {
|
||||
node_type: "KSampler",
|
||||
title: "KSampler",
|
||||
category: "sampling",
|
||||
description: "Samples latent images with a diffusion model.",
|
||||
inputs: json!({
|
||||
"model": { "type": "MODEL", "required": true },
|
||||
"positive": { "type": "CONDITIONING", "required": true },
|
||||
"negative": { "type": "CONDITIONING", "required": true },
|
||||
"latent_image": { "type": "LATENT", "required": true }
|
||||
}),
|
||||
params: json!({
|
||||
"seed": { "ui": "number", "type": "INT", "default": 0 },
|
||||
"control_after_generate": { "ui": "boolean", "type": "BOOLEAN", "default": true },
|
||||
"steps": { "ui": "number", "type": "INT", "default": 20, "min": 1, "max": 100 },
|
||||
"cfg": { "ui": "number", "type": "FLOAT", "default": 8, "min": 0, "max": 50 },
|
||||
"sampler_name": { "ui": "text", "type": "STRING", "default": "euler" },
|
||||
"scheduler": { "ui": "text", "type": "STRING", "default": "normal" },
|
||||
"denoise": { "ui": "number", "type": "FLOAT", "default": 1, "min": 0, "max": 1 }
|
||||
}),
|
||||
outputs: json!({ "LATENT": { "type": "LATENT" } }),
|
||||
},
|
||||
NodeTypeDefinition {
|
||||
node_type: "VAEDecode",
|
||||
title: "VAEDecode",
|
||||
category: "latent",
|
||||
description: "Decodes latent samples to images.",
|
||||
inputs: json!({
|
||||
"samples": { "type": "LATENT", "required": true },
|
||||
"vae": { "type": "VAE", "required": true }
|
||||
}),
|
||||
params: json!({}),
|
||||
outputs: json!({ "IMAGE": { "type": "IMAGE" } }),
|
||||
},
|
||||
NodeTypeDefinition {
|
||||
node_type: "SaveImage",
|
||||
title: "SaveImage",
|
||||
category: "image",
|
||||
description: "Saves images to disk.",
|
||||
inputs: json!({
|
||||
"images": { "type": "IMAGE", "required": true }
|
||||
}),
|
||||
params: json!({
|
||||
"filename_prefix": { "ui": "text", "type": "STRING", "default": "ComfyUI" }
|
||||
}),
|
||||
outputs: json!({ "IMAGE": { "type": "IMAGE" } }),
|
||||
},
|
||||
NodeTypeDefinition {
|
||||
node_type: "PreviewImage",
|
||||
title: "PreviewImage",
|
||||
category: "image",
|
||||
description: "Previews generated images.",
|
||||
inputs: json!({
|
||||
"images": { "type": "IMAGE", "required": true }
|
||||
}),
|
||||
params: json!({}),
|
||||
outputs: json!({ "IMAGE": { "type": "IMAGE" } }),
|
||||
},
|
||||
]
|
||||
}
|
||||
|
||||
pub fn node_types() -> Vec<NodeTypeDefinition> {
|
||||
let mut nodes = core_node_types();
|
||||
nodes.extend(vec![
|
||||
NodeTypeDefinition {
|
||||
node_type: "input.text",
|
||||
title: "文本输入",
|
||||
category: "Input",
|
||||
description: "输入脚本、提示词或任意文本。",
|
||||
inputs: json!({}),
|
||||
params: json!({
|
||||
"value": { "ui": "textarea", "type": "STRING", "default": "一个电影感的雨夜城市镜头" }
|
||||
}),
|
||||
outputs: json!({ "text": { "type": "STRING" } }),
|
||||
},
|
||||
NodeTypeDefinition {
|
||||
node_type: "boosapi.comfyui_workflow",
|
||||
title: "生成 ComfyUI 工作流",
|
||||
category: "ComfyUI",
|
||||
description: "用 BoosAPI 模型把文本需求改写为 ComfyUI workflow JSON。",
|
||||
inputs: json!({ "text": { "type": "STRING", "required": true } }),
|
||||
params: json!({
|
||||
"model": { "ui": "text", "type": "STRING", "default": "gpt-5.4" },
|
||||
"system": { "ui": "textarea", "type": "STRING", "default": "You generate valid ComfyUI workflow JSON. Return JSON only." },
|
||||
"temperature": { "ui": "number", "type": "NUMBER", "default": 0.3, "min": 0, "max": 2 }
|
||||
}),
|
||||
outputs: json!({ "workflow": { "type": "JSON" }, "text": { "type": "STRING" }, "raw": { "type": "JSON" } }),
|
||||
},
|
||||
NodeTypeDefinition {
|
||||
node_type: "comfyui.submit_workflow",
|
||||
title: "提交到外部 ComfyUI",
|
||||
category: "ComfyUI",
|
||||
description: "把 workflow JSON 提交到已配置的外部 ComfyUI /prompt。",
|
||||
inputs: json!({ "workflow": { "type": "JSON", "required": true } }),
|
||||
params: json!({
|
||||
"client_id": { "ui": "text", "type": "STRING", "default": "boosapi-workflow-studio" },
|
||||
"dry_run": { "ui": "boolean", "type": "BOOLEAN", "default": true }
|
||||
}),
|
||||
outputs: json!({ "prompt_id": { "type": "STRING" }, "response": { "type": "JSON" }, "workflow": { "type": "JSON" } }),
|
||||
},
|
||||
NodeTypeDefinition {
|
||||
node_type: "boosapi.chat_completion",
|
||||
title: "BoosAPI Chat Completion",
|
||||
category: "BoosAPI/Text",
|
||||
description: "调用 BoosAPI OpenAI 兼容聊天模型。",
|
||||
inputs: json!({ "text": { "type": "STRING", "required": true } }),
|
||||
params: json!({
|
||||
"model": { "ui": "text", "type": "STRING", "default": "gpt-5.4" },
|
||||
"system": { "ui": "textarea", "type": "STRING", "default": "" },
|
||||
"temperature": { "ui": "number", "type": "NUMBER", "default": 0.7, "min": 0, "max": 2 }
|
||||
}),
|
||||
outputs: json!({ "text": { "type": "STRING" }, "raw": { "type": "JSON" } }),
|
||||
},
|
||||
NodeTypeDefinition {
|
||||
node_type: "boosapi.text_to_image",
|
||||
title: "文生图",
|
||||
category: "BoosAPI/Image",
|
||||
description: "通过 /v1/images/generations 生成图片，并归档到工作空间。",
|
||||
inputs: json!({ "prompt": { "type": "STRING", "required": true } }),
|
||||
params: json!({
|
||||
"model": { "ui": "text", "type": "STRING", "default": "gpt-image-2" },
|
||||
"width": { "ui": "number", "type": "NUMBER", "default": 1024, "min": 64, "max": 4096 },
|
||||
"height": { "ui": "number", "type": "NUMBER", "default": 1024, "min": 64, "max": 4096 },
|
||||
"size": { "ui": "text", "type": "STRING", "default": "" }
|
||||
}),
|
||||
outputs: json!({
|
||||
"image": { "type": "IMAGE" },
|
||||
"images": { "type": "IMAGE", "is_list": true },
|
||||
"raw": { "type": "JSON" }
|
||||
}),
|
||||
},
|
||||
NodeTypeDefinition {
|
||||
node_type: "boosapi.text_to_video",
|
||||
title: "文生视频",
|
||||
category: "BoosAPI/Video",
|
||||
description: "通过 /v1/videos 创建文生视频任务。",
|
||||
inputs: json!({ "prompt": { "type": "STRING", "required": true } }),
|
||||
params: json!({
|
||||
"model": { "ui": "text", "type": "STRING", "default": "seedance-1-0-lite-t2v" },
|
||||
"ratio": { "ui": "text", "type": "STRING", "default": "16:9" },
|
||||
"duration": { "ui": "number", "type": "NUMBER", "default": 5, "min": 1, "max": 30 },
|
||||
"resolution": { "ui": "text", "type": "STRING", "default": "720p" },
|
||||
"fps": { "ui": "number", "type": "NUMBER", "default": 24, "min": 1, "max": 60 }
|
||||
}),
|
||||
outputs: json!({
|
||||
"task_id": { "type": "STRING" },
|
||||
"video": { "type": "VIDEO" },
|
||||
"raw": { "type": "JSON" }
|
||||
}),
|
||||
},
|
||||
NodeTypeDefinition {
|
||||
node_type: "boosapi.image_to_video",
|
||||
title: "图生视频",
|
||||
category: "BoosAPI/Video",
|
||||
description: "通过 /v1/videos 创建图生视频任务。",
|
||||
inputs: json!({
|
||||
"prompt": { "type": "STRING", "required": true },
|
||||
"image": { "type": "IMAGE", "required": false }
|
||||
}),
|
||||
params: json!({
|
||||
"model": { "ui": "text", "type": "STRING", "default": "seedance-1-0-lite-i2v" },
|
||||
"image_url": { "ui": "text", "type": "STRING", "default": "" },
|
||||
"ratio": { "ui": "text", "type": "STRING", "default": "16:9" },
|
||||
"duration": { "ui": "number", "type": "NUMBER", "default": 5, "min": 1, "max": 30 },
|
||||
"resolution": { "ui": "text", "type": "STRING", "default": "720p" },
|
||||
"fps": { "ui": "number", "type": "NUMBER", "default": 24, "min": 1, "max": 60 }
|
||||
}),
|
||||
outputs: json!({
|
||||
"task_id": { "type": "STRING" },
|
||||
"video": { "type": "VIDEO" },
|
||||
"raw": { "type": "JSON" }
|
||||
}),
|
||||
},
|
||||
NodeTypeDefinition {
|
||||
node_type: "boosapi.video_status",
|
||||
title: "查询视频任务",
|
||||
category: "BoosAPI/Video",
|
||||
description: "查询 /v1/videos/{id} 视频任务结果。",
|
||||
inputs: json!({ "task_id": { "type": "STRING", "required": true } }),
|
||||
params: json!({
|
||||
"model": { "ui": "text", "type": "STRING", "default": "seedance-1-0-lite-t2v" }
|
||||
}),
|
||||
outputs: json!({
|
||||
"video": { "type": "VIDEO" },
|
||||
"raw": { "type": "JSON" }
|
||||
}),
|
||||
},
|
||||
NodeTypeDefinition {
|
||||
node_type: "output.preview_image",
|
||||
title: "图片预览",
|
||||
category: "Output",
|
||||
description: "收集图片输出，供运行结果面板展示。",
|
||||
inputs: json!({ "image": { "type": "IMAGE", "required": true } }),
|
||||
params: json!({}),
|
||||
outputs: json!({ "image": { "type": "IMAGE" } }),
|
||||
},
|
||||
NodeTypeDefinition {
|
||||
node_type: "output.preview_video",
|
||||
title: "视频预览",
|
||||
category: "Output",
|
||||
description: "收集视频输出，供运行结果面板展示。",
|
||||
inputs: json!({ "video": { "type": "VIDEO", "required": true } }),
|
||||
params: json!({}),
|
||||
outputs: json!({ "video": { "type": "VIDEO" } }),
|
||||
},
|
||||
NodeTypeDefinition {
|
||||
node_type: "output.save_file",
|
||||
title: "保存文件",
|
||||
category: "Output",
|
||||
description: "标记上游文件为最终输出。",
|
||||
inputs: json!({ "file": { "type": "FILE", "required": true } }),
|
||||
params: json!({ "name": { "ui": "text", "type": "STRING", "default": "" } }),
|
||||
outputs: json!({ "file": { "type": "FILE" } }),
|
||||
},
|
||||
NodeTypeDefinition {
|
||||
node_type: "util.text_join",
|
||||
title: "文本拼接",
|
||||
category: "Util",
|
||||
description: "拼接输入文本和静态文本。",
|
||||
inputs: json!({ "text": { "type": "STRING", "required": false } }),
|
||||
params: json!({
|
||||
"prefix": { "ui": "textarea", "type": "STRING", "default": "" },
|
||||
"suffix": { "ui": "textarea", "type": "STRING", "default": "" },
|
||||
"separator": { "ui": "text", "type": "STRING", "default": "\\n" }
|
||||
}),
|
||||
outputs: json!({ "text": { "type": "STRING" } }),
|
||||
},
|
||||
NodeTypeDefinition {
|
||||
node_type: "util.json_extract",
|
||||
title: "JSON 提取",
|
||||
category: "Util",
|
||||
description: "从 JSON 对象中按点路径提取字段。",
|
||||
inputs: json!({ "json": { "type": "JSON", "required": true } }),
|
||||
params: json!({ "path": { "ui": "text", "type": "STRING", "default": "data.0.url" } }),
|
||||
outputs: json!({ "value": { "type": "JSON" }, "text": { "type": "STRING" } }),
|
||||
},
|
||||
]);
|
||||
nodes
|
||||
}
|
||||
|
||||
pub fn is_known_node(node_type: &str) -> bool {
|
||||
node_types()
|
||||
.iter()
|
||||
.any(|definition| definition.node_type == node_type)
|
||||
}
|
||||
@@ -0,0 +1,389 @@
|
||||
use crate::{
|
||||
auth::jwt::validate_jwt,
|
||||
db::models::downstream_api_key::DownstreamApiKey,
|
||||
error::{AppError, AppResult},
|
||||
workflow::{
|
||||
executor, graph,
|
||||
models::{default_graph, RunContext, WorkflowGraph},
|
||||
registry, storage,
|
||||
},
|
||||
AppState,
|
||||
};
|
||||
use axum::{
|
||||
extract::{Path, Query, State},
|
||||
http::{header, HeaderMap, StatusCode},
|
||||
routing::{get, post},
|
||||
Json, Router,
|
||||
};
|
||||
use serde::Deserialize;
|
||||
use serde_json::{json, Value};
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
enum WorkflowCaller {
|
||||
Admin,
|
||||
User { user_id: i64 },
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct WorkflowListQuery {
|
||||
#[serde(default)]
|
||||
include_public: bool,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct RunListQuery {
|
||||
limit: Option<i64>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
struct SaveWorkflowPayload {
|
||||
name: Option<String>,
|
||||
description: Option<String>,
|
||||
graph: Option<Value>,
|
||||
visibility: Option<String>,
|
||||
is_template: Option<bool>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct RunWorkflowPayload {
|
||||
inputs: Option<Value>,
|
||||
}
|
||||
|
||||
pub fn routes() -> Router<AppState> {
|
||||
Router::new()
|
||||
.route("/api/workflow/node-types", get(node_types))
|
||||
.route("/api/workflows", get(list_workflows).post(create_workflow))
|
||||
.route(
|
||||
"/api/workflows/{id}",
|
||||
get(get_workflow)
|
||||
.put(update_workflow)
|
||||
.delete(delete_workflow),
|
||||
)
|
||||
.route("/api/workflows/{id}/run", post(run_workflow))
|
||||
.route("/api/workflow-runs", get(list_runs))
|
||||
.route("/api/workflow-runs/{id}", get(get_run))
|
||||
.route("/api/workflow-runs/{id}/outputs", get(get_run_outputs))
|
||||
.route("/api/user/workflow/node-types", get(node_types))
|
||||
.route(
|
||||
"/api/user/workflows",
|
||||
get(list_workflows).post(create_workflow),
|
||||
)
|
||||
.route(
|
||||
"/api/user/workflows/{id}",
|
||||
get(get_workflow)
|
||||
.put(update_workflow)
|
||||
.delete(delete_workflow),
|
||||
)
|
||||
.route("/api/user/workflows/{id}/run", post(run_workflow))
|
||||
.route("/api/user/workflow-runs", get(list_runs))
|
||||
.route("/api/user/workflow-runs/{id}", get(get_run))
|
||||
.route("/api/user/workflow-runs/{id}/outputs", get(get_run_outputs))
|
||||
}
|
||||
|
||||
fn bearer(headers: &HeaderMap) -> Option<&str> {
|
||||
headers
|
||||
.get(header::AUTHORIZATION)
|
||||
.and_then(|value| value.to_str().ok())
|
||||
.and_then(|value| {
|
||||
value
|
||||
.strip_prefix("Bearer ")
|
||||
.or_else(|| value.strip_prefix("bearer "))
|
||||
})
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
}
|
||||
|
||||
fn resolve_caller(headers: &HeaderMap, state: &AppState) -> AppResult<WorkflowCaller> {
|
||||
let token = bearer(headers)
|
||||
.ok_or_else(|| AppError::Unauthorized("Missing Authorization header".into()))?;
|
||||
if token == state.config.auth.auth_token || token == state.config.auth.proxy_token {
|
||||
return Ok(WorkflowCaller::Admin);
|
||||
}
|
||||
let claims = validate_jwt(&state.config.auth.jwt_secret, token)
|
||||
.map_err(|_| AppError::Unauthorized("Invalid or expired token".into()))?;
|
||||
if claims.role == "admin" {
|
||||
return Ok(WorkflowCaller::Admin);
|
||||
}
|
||||
if claims.role == "user" {
|
||||
let user_id = claims
|
||||
.sub
|
||||
.parse::<i64>()
|
||||
.map_err(|_| AppError::Unauthorized("Invalid user ID in token".into()))?;
|
||||
return Ok(WorkflowCaller::User { user_id });
|
||||
}
|
||||
Err(AppError::Forbidden("Workflow access denied".into()))
|
||||
}
|
||||
|
||||
fn caller_user_id(caller: &WorkflowCaller) -> Option<i64> {
|
||||
match caller {
|
||||
WorkflowCaller::Admin => None,
|
||||
WorkflowCaller::User { user_id } => Some(*user_id),
|
||||
}
|
||||
}
|
||||
|
||||
fn caller_is_admin(caller: &WorkflowCaller) -> bool {
|
||||
matches!(caller, WorkflowCaller::Admin)
|
||||
}
|
||||
|
||||
async fn user_downstream_key(state: &AppState, user_id: i64) -> AppResult<DownstreamApiKey> {
|
||||
sqlx::query_as::<_, DownstreamApiKey>(
|
||||
"SELECT *
|
||||
FROM downstream_api_keys
|
||||
WHERE user_id = ?
|
||||
AND COALESCE(enabled, 1) = 1
|
||||
AND (
|
||||
expires_at IS NULL
|
||||
OR expires_at = ''
|
||||
OR datetime(expires_at) > datetime('now')
|
||||
)
|
||||
ORDER BY updated_at DESC, created_at DESC, id DESC
|
||||
LIMIT 1",
|
||||
)
|
||||
.bind(user_id)
|
||||
.fetch_optional(&*state.db)
|
||||
.await?
|
||||
.ok_or_else(|| {
|
||||
AppError::Forbidden("当前用户没有可用的 API Key，请先在用户中心创建或启用 API Key".into())
|
||||
})
|
||||
}
|
||||
|
||||
async fn run_context(state: &AppState, caller: &WorkflowCaller) -> AppResult<RunContext> {
|
||||
let bearer = match caller {
|
||||
WorkflowCaller::Admin => state.builtin_keys.comfyui_agent.1.clone(),
|
||||
WorkflowCaller::User { user_id } => user_downstream_key(state, *user_id).await?.key,
|
||||
};
|
||||
let host = if state.config.server.host == "0.0.0.0" || state.config.server.host == "::" {
|
||||
"127.0.0.1".to_string()
|
||||
} else {
|
||||
state.config.server.host.clone()
|
||||
};
|
||||
Ok(RunContext {
|
||||
user_id: caller_user_id(caller),
|
||||
bearer,
|
||||
base_url: format!("http://{}:{}/v1", host, state.config.server.port),
|
||||
})
|
||||
}
|
||||
|
||||
fn normalize_graph(value: Option<Value>, name: &str) -> AppResult<Value> {
|
||||
let graph = value.unwrap_or_else(|| default_graph(name));
|
||||
let parsed: WorkflowGraph = serde_json::from_value(graph.clone())
|
||||
.map_err(|err| AppError::BadRequest(format!("invalid workflow graph: {err}")))?;
|
||||
graph::validate_and_sort(&parsed)?;
|
||||
Ok(graph)
|
||||
}
|
||||
|
||||
fn normalize_visibility(raw: Option<String>, caller: &WorkflowCaller) -> String {
|
||||
let value = raw.unwrap_or_else(|| "private".into());
|
||||
match value.as_str() {
|
||||
"public" | "template" if caller_is_admin(caller) => value,
|
||||
_ => "private".into(),
|
||||
}
|
||||
}
|
||||
|
||||
async fn node_types() -> Json<Value> {
|
||||
Json(json!({ "success": true, "nodeTypes": registry::node_types() }))
|
||||
}
|
||||
|
||||
async fn list_workflows(
|
||||
State(state): State<AppState>,
|
||||
headers: HeaderMap,
|
||||
Query(query): Query<WorkflowListQuery>,
|
||||
) -> AppResult<Json<Value>> {
|
||||
let caller = resolve_caller(&headers, &state)?;
|
||||
let items = storage::list_workflows(
|
||||
&state.db,
|
||||
caller_user_id(&caller),
|
||||
query.include_public || caller_is_admin(&caller),
|
||||
)
|
||||
.await?;
|
||||
Ok(Json(json!({ "success": true, "items": items })))
|
||||
}
|
||||
|
||||
async fn create_workflow(
|
||||
State(state): State<AppState>,
|
||||
headers: HeaderMap,
|
||||
Json(payload): Json<SaveWorkflowPayload>,
|
||||
) -> AppResult<(StatusCode, Json<Value>)> {
|
||||
let caller = resolve_caller(&headers, &state)?;
|
||||
let name = payload
|
||||
.name
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.unwrap_or("未命名工作流");
|
||||
let graph = normalize_graph(payload.graph, name)?;
|
||||
let visibility = normalize_visibility(payload.visibility, &caller);
|
||||
let workflow = storage::insert_workflow(
|
||||
&state,
|
||||
caller_user_id(&caller),
|
||||
name,
|
||||
payload.description.as_deref(),
|
||||
&graph,
|
||||
&visibility,
|
||||
payload.is_template.unwrap_or(false) && caller_is_admin(&caller),
|
||||
)
|
||||
.await?;
|
||||
Ok((
|
||||
StatusCode::CREATED,
|
||||
Json(json!({ "success": true, "workflow": workflow })),
|
||||
))
|
||||
}
|
||||
|
||||
async fn get_workflow(
|
||||
State(state): State<AppState>,
|
||||
headers: HeaderMap,
|
||||
Path(id): Path<String>,
|
||||
) -> AppResult<Json<Value>> {
|
||||
let caller = resolve_caller(&headers, &state)?;
|
||||
let workflow = storage::get_workflow(
|
||||
&state.db,
|
||||
&id,
|
||||
caller_user_id(&caller),
|
||||
caller_is_admin(&caller),
|
||||
)
|
||||
.await?;
|
||||
Ok(Json(json!({ "success": true, "workflow": workflow })))
|
||||
}
|
||||
|
||||
async fn update_workflow(
|
||||
State(state): State<AppState>,
|
||||
headers: HeaderMap,
|
||||
Path(id): Path<String>,
|
||||
Json(payload): Json<SaveWorkflowPayload>,
|
||||
) -> AppResult<Json<Value>> {
|
||||
let caller = resolve_caller(&headers, &state)?;
|
||||
let current = storage::get_workflow(
|
||||
&state.db,
|
||||
&id,
|
||||
caller_user_id(&caller),
|
||||
caller_is_admin(&caller),
|
||||
)
|
||||
.await?;
|
||||
let name = payload
|
||||
.name
|
||||
.as_deref()
|
||||
.map(str::trim)
|
||||
.filter(|value| !value.is_empty())
|
||||
.unwrap_or(¤t.name);
|
||||
let graph = normalize_graph(payload.graph.or(Some(current.graph)), name)?;
|
||||
let visibility = normalize_visibility(payload.visibility.or(Some(current.visibility)), &caller);
|
||||
let workflow = storage::update_workflow(
|
||||
&state,
|
||||
&id,
|
||||
caller_user_id(&caller),
|
||||
caller_is_admin(&caller),
|
||||
name,
|
||||
payload
|
||||
.description
|
||||
.as_deref()
|
||||
.or(current.description.as_deref()),
|
||||
&graph,
|
||||
&visibility,
|
||||
)
|
||||
.await?;
|
||||
Ok(Json(json!({ "success": true, "workflow": workflow })))
|
||||
}
|
||||
|
||||
async fn delete_workflow(
|
||||
State(state): State<AppState>,
|
||||
headers: HeaderMap,
|
||||
Path(id): Path<String>,
|
||||
) -> AppResult<StatusCode> {
|
||||
let caller = resolve_caller(&headers, &state)?;
|
||||
storage::delete_workflow(
|
||||
&state,
|
||||
&id,
|
||||
caller_user_id(&caller),
|
||||
caller_is_admin(&caller),
|
||||
)
|
||||
.await?;
|
||||
Ok(StatusCode::NO_CONTENT)
|
||||
}
|
||||
|
||||
async fn run_workflow(
|
||||
State(state): State<AppState>,
|
||||
headers: HeaderMap,
|
||||
Path(id): Path<String>,
|
||||
Json(payload): Json<RunWorkflowPayload>,
|
||||
) -> AppResult<Json<Value>> {
|
||||
let caller = resolve_caller(&headers, &state)?;
|
||||
let workflow = storage::get_workflow(
|
||||
&state.db,
|
||||
&id,
|
||||
caller_user_id(&caller),
|
||||
caller_is_admin(&caller),
|
||||
)
|
||||
.await?;
|
||||
let ctx = run_context(&state, &caller).await?;
|
||||
let run_id = executor::run_workflow(
|
||||
&state,
|
||||
workflow.id,
|
||||
workflow.graph,
|
||||
payload.inputs.unwrap_or_else(|| json!({})),
|
||||
ctx,
|
||||
)
|
||||
.await?;
|
||||
let run = storage::get_run(
|
||||
&state.db,
|
||||
&run_id,
|
||||
caller_user_id(&caller),
|
||||
caller_is_admin(&caller),
|
||||
)
|
||||
.await?;
|
||||
Ok(Json(
|
||||
json!({ "success": true, "runId": run_id, "run": run }),
|
||||
))
|
||||
}
|
||||
|
||||
async fn list_runs(
|
||||
State(state): State<AppState>,
|
||||
headers: HeaderMap,
|
||||
Query(query): Query<RunListQuery>,
|
||||
) -> AppResult<Json<Value>> {
|
||||
let caller = resolve_caller(&headers, &state)?;
|
||||
let items = storage::list_runs(
|
||||
&state.db,
|
||||
caller_user_id(&caller),
|
||||
caller_is_admin(&caller),
|
||||
query.limit.unwrap_or(30),
|
||||
)
|
||||
.await?;
|
||||
Ok(Json(json!({ "success": true, "items": items })))
|
||||
}
|
||||
|
||||
async fn get_run(
|
||||
State(state): State<AppState>,
|
||||
headers: HeaderMap,
|
||||
Path(id): Path<String>,
|
||||
) -> AppResult<Json<Value>> {
|
||||
let caller = resolve_caller(&headers, &state)?;
|
||||
let run = storage::get_run(
|
||||
&state.db,
|
||||
&id,
|
||||
caller_user_id(&caller),
|
||||
caller_is_admin(&caller),
|
||||
)
|
||||
.await?;
|
||||
Ok(Json(json!({ "success": true, "run": run })))
|
||||
}
|
||||
|
||||
async fn get_run_outputs(
|
||||
State(state): State<AppState>,
|
||||
headers: HeaderMap,
|
||||
Path(id): Path<String>,
|
||||
) -> AppResult<Json<Value>> {
|
||||
let caller = resolve_caller(&headers, &state)?;
|
||||
let run = storage::get_run(
|
||||
&state.db,
|
||||
&id,
|
||||
caller_user_id(&caller),
|
||||
caller_is_admin(&caller),
|
||||
)
|
||||
.await?;
|
||||
Ok(Json(json!({
|
||||
"success": true,
|
||||
"runId": run.public_id,
|
||||
"outputs": run.final_outputs.unwrap_or_else(|| json!({}))
|
||||
})))
|
||||
}
|
||||
@@ -0,0 +1,463 @@
|
||||
use crate::{
|
||||
error::{AppError, AppResult},
|
||||
workflow::models::{WorkflowNodeRunSummary, WorkflowRunSummary, WorkflowSummary},
|
||||
AppState,
|
||||
};
|
||||
use serde_json::{json, Value};
|
||||
use sqlx::{Row, SqlitePool};
|
||||
use uuid::Uuid;
|
||||
|
||||
pub fn new_workflow_id() -> String {
|
||||
format!("wf_{}", Uuid::new_v4().simple())
|
||||
}
|
||||
|
||||
pub fn new_run_id() -> String {
|
||||
format!("run_{}", Uuid::new_v4().simple())
|
||||
}
|
||||
|
||||
pub fn parse_json_text(raw: Option<String>, fallback: Value) -> Value {
|
||||
raw.and_then(|value| serde_json::from_str::<Value>(&value).ok())
|
||||
.unwrap_or(fallback)
|
||||
}
|
||||
|
||||
fn row_to_workflow(row: sqlx::sqlite::SqliteRow) -> WorkflowSummary {
|
||||
let graph_raw: String = row.try_get("graph").unwrap_or_else(|_| "{}".to_string());
|
||||
WorkflowSummary {
|
||||
id: row.try_get("id").unwrap_or_default(),
|
||||
public_id: row.try_get("public_id").unwrap_or_default(),
|
||||
user_id: row.try_get::<Option<i64>, _>("user_id").ok().flatten(),
|
||||
name: row.try_get("name").unwrap_or_default(),
|
||||
description: row.try_get("description").ok(),
|
||||
graph: serde_json::from_str(&graph_raw).unwrap_or_else(|_| json!({})),
|
||||
is_template: row
|
||||
.try_get::<i64, _>("is_template")
|
||||
.map(|value| value != 0)
|
||||
.unwrap_or(false),
|
||||
visibility: row
|
||||
.try_get("visibility")
|
||||
.unwrap_or_else(|_| "private".into()),
|
||||
created_at: row.try_get("created_at").unwrap_or_default(),
|
||||
updated_at: row.try_get("updated_at").unwrap_or_default(),
|
||||
}
|
||||
}
|
||||
|
||||
fn row_to_node_run(row: sqlx::sqlite::SqliteRow) -> WorkflowNodeRunSummary {
|
||||
WorkflowNodeRunSummary {
|
||||
id: row.try_get("id").unwrap_or_default(),
|
||||
node_id: row.try_get("node_id").unwrap_or_default(),
|
||||
node_type: row.try_get("node_type").unwrap_or_default(),
|
||||
status: row.try_get("status").unwrap_or_default(),
|
||||
inputs: parse_json_text(row.try_get("inputs").ok(), json!({})),
|
||||
outputs: row
|
||||
.try_get::<Option<String>, _>("outputs")
|
||||
.ok()
|
||||
.flatten()
|
||||
.and_then(|raw| serde_json::from_str(&raw).ok()),
|
||||
error_message: row.try_get("error_message").ok(),
|
||||
started_at: row.try_get("started_at").ok(),
|
||||
completed_at: row.try_get("completed_at").ok(),
|
||||
created_at: row.try_get("created_at").unwrap_or_default(),
|
||||
}
|
||||
}
|
||||
|
||||
fn row_to_run(
|
||||
row: sqlx::sqlite::SqliteRow,
|
||||
node_runs: Vec<WorkflowNodeRunSummary>,
|
||||
) -> WorkflowRunSummary {
|
||||
WorkflowRunSummary {
|
||||
id: row.try_get("id").unwrap_or_default(),
|
||||
public_id: row.try_get("public_id").unwrap_or_default(),
|
||||
workflow_id: row.try_get("workflow_id").unwrap_or_default(),
|
||||
workflow_public_id: row.try_get("workflow_public_id").ok(),
|
||||
user_id: row.try_get::<Option<i64>, _>("user_id").ok().flatten(),
|
||||
status: row.try_get("status").unwrap_or_default(),
|
||||
input_overrides: parse_json_text(row.try_get("input_overrides").ok(), json!({})),
|
||||
final_outputs: row
|
||||
.try_get::<Option<String>, _>("final_outputs")
|
||||
.ok()
|
||||
.flatten()
|
||||
.and_then(|raw| serde_json::from_str(&raw).ok()),
|
||||
error_message: row.try_get("error_message").ok(),
|
||||
started_at: row.try_get("started_at").ok(),
|
||||
completed_at: row.try_get("completed_at").ok(),
|
||||
created_at: row.try_get("created_at").unwrap_or_default(),
|
||||
node_runs,
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn list_workflows(
|
||||
pool: &SqlitePool,
|
||||
user_id: Option<i64>,
|
||||
include_public: bool,
|
||||
) -> AppResult<Vec<WorkflowSummary>> {
|
||||
let rows = sqlx::query(
|
||||
"SELECT *
|
||||
FROM workflows
|
||||
WHERE deleted_at IS NULL
|
||||
AND (
|
||||
? IS NULL
|
||||
OR user_id = ?
|
||||
OR (? = 1 AND visibility IN ('public', 'template'))
|
||||
)
|
||||
ORDER BY datetime(updated_at) DESC, id DESC",
|
||||
)
|
||||
.bind(user_id)
|
||||
.bind(user_id)
|
||||
.bind(if include_public { 1_i64 } else { 0_i64 })
|
||||
.map(row_to_workflow)
|
||||
.fetch_all(pool)
|
||||
.await?;
|
||||
Ok(rows)
|
||||
}
|
||||
|
||||
pub async fn get_workflow(
|
||||
pool: &SqlitePool,
|
||||
public_id: &str,
|
||||
user_id: Option<i64>,
|
||||
allow_admin: bool,
|
||||
) -> AppResult<WorkflowSummary> {
|
||||
let row = sqlx::query(
|
||||
"SELECT *
|
||||
FROM workflows
|
||||
WHERE public_id = ?
|
||||
AND deleted_at IS NULL
|
||||
AND (? = 1 OR user_id = ? OR visibility IN ('public', 'template'))
|
||||
LIMIT 1",
|
||||
)
|
||||
.bind(public_id)
|
||||
.bind(if allow_admin { 1_i64 } else { 0_i64 })
|
||||
.bind(user_id)
|
||||
.fetch_optional(pool)
|
||||
.await?
|
||||
.ok_or_else(|| AppError::NotFound("Workflow not found".into()))?;
|
||||
Ok(row_to_workflow(row))
|
||||
}
|
||||
|
||||
pub async fn get_workflow_by_db_id(pool: &SqlitePool, id: i64) -> AppResult<WorkflowSummary> {
|
||||
let row = sqlx::query("SELECT * FROM workflows WHERE id = ? AND deleted_at IS NULL LIMIT 1")
|
||||
.bind(id)
|
||||
.fetch_optional(pool)
|
||||
.await?
|
||||
.ok_or_else(|| AppError::NotFound("Workflow not found".into()))?;
|
||||
Ok(row_to_workflow(row))
|
||||
}
|
||||
|
||||
pub async fn get_run(
|
||||
pool: &SqlitePool,
|
||||
public_id: &str,
|
||||
user_id: Option<i64>,
|
||||
allow_admin: bool,
|
||||
) -> AppResult<WorkflowRunSummary> {
|
||||
let row = sqlx::query(
|
||||
"SELECT workflow_runs.*, workflows.public_id AS workflow_public_id
|
||||
FROM workflow_runs
|
||||
JOIN workflows ON workflows.id = workflow_runs.workflow_id
|
||||
WHERE workflow_runs.public_id = ?
|
||||
AND (? = 1 OR workflow_runs.user_id = ?)
|
||||
LIMIT 1",
|
||||
)
|
||||
.bind(public_id)
|
||||
.bind(if allow_admin { 1_i64 } else { 0_i64 })
|
||||
.bind(user_id)
|
||||
.fetch_optional(pool)
|
||||
.await?
|
||||
.ok_or_else(|| AppError::NotFound("Workflow run not found".into()))?;
|
||||
let run_id: i64 = row.try_get("id").unwrap_or_default();
|
||||
let node_runs = sqlx::query(
|
||||
"SELECT *
|
||||
FROM workflow_node_runs
|
||||
WHERE run_id = ?
|
||||
ORDER BY id ASC",
|
||||
)
|
||||
.bind(run_id)
|
||||
.map(row_to_node_run)
|
||||
.fetch_all(pool)
|
||||
.await?;
|
||||
Ok(row_to_run(row, node_runs))
|
||||
}
|
||||
|
||||
pub async fn list_runs(
|
||||
pool: &SqlitePool,
|
||||
user_id: Option<i64>,
|
||||
allow_admin: bool,
|
||||
limit: i64,
|
||||
) -> AppResult<Vec<WorkflowRunSummary>> {
|
||||
let rows = sqlx::query(
|
||||
"SELECT workflow_runs.*, workflows.public_id AS workflow_public_id
|
||||
FROM workflow_runs
|
||||
JOIN workflows ON workflows.id = workflow_runs.workflow_id
|
||||
WHERE (? = 1 OR workflow_runs.user_id = ?)
|
||||
ORDER BY datetime(workflow_runs.created_at) DESC, workflow_runs.id DESC
|
||||
LIMIT ?",
|
||||
)
|
||||
.bind(if allow_admin { 1_i64 } else { 0_i64 })
|
||||
.bind(user_id)
|
||||
.bind(limit.clamp(1, 100))
|
||||
.fetch_all(pool)
|
||||
.await?;
|
||||
|
||||
let mut runs = Vec::with_capacity(rows.len());
|
||||
for row in rows {
|
||||
let run_id: i64 = row.try_get("id").unwrap_or_default();
|
||||
let node_runs = sqlx::query(
|
||||
"SELECT *
|
||||
FROM workflow_node_runs
|
||||
WHERE run_id = ?
|
||||
ORDER BY id ASC",
|
||||
)
|
||||
.bind(run_id)
|
||||
.map(row_to_node_run)
|
||||
.fetch_all(pool)
|
||||
.await?;
|
||||
runs.push(row_to_run(row, node_runs));
|
||||
}
|
||||
Ok(runs)
|
||||
}
|
||||
|
||||
pub async fn insert_workflow(
|
||||
state: &AppState,
|
||||
user_id: Option<i64>,
|
||||
name: &str,
|
||||
description: Option<&str>,
|
||||
graph: &Value,
|
||||
visibility: &str,
|
||||
is_template: bool,
|
||||
) -> AppResult<WorkflowSummary> {
|
||||
let public_id = new_workflow_id();
|
||||
sqlx::query(
|
||||
"INSERT INTO workflows
|
||||
(public_id, user_id, name, description, graph, visibility, is_template)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?)",
|
||||
)
|
||||
.bind(&public_id)
|
||||
.bind(user_id)
|
||||
.bind(name)
|
||||
.bind(description)
|
||||
.bind(graph.to_string())
|
||||
.bind(visibility)
|
||||
.bind(if is_template { 1_i64 } else { 0_i64 })
|
||||
.execute(&*state.db)
|
||||
.await?;
|
||||
get_workflow(&state.db, &public_id, user_id, true).await
|
||||
}
|
||||
|
||||
pub async fn update_workflow(
|
||||
state: &AppState,
|
||||
public_id: &str,
|
||||
user_id: Option<i64>,
|
||||
allow_admin: bool,
|
||||
name: &str,
|
||||
description: Option<&str>,
|
||||
graph: &Value,
|
||||
visibility: &str,
|
||||
) -> AppResult<WorkflowSummary> {
|
||||
let result = sqlx::query(
|
||||
"UPDATE workflows
|
||||
SET name = ?, description = ?, graph = ?, visibility = ?, updated_at = datetime('now')
|
||||
WHERE public_id = ?
|
||||
AND deleted_at IS NULL
|
||||
AND (? = 1 OR user_id = ?)",
|
||||
)
|
||||
.bind(name)
|
||||
.bind(description)
|
||||
.bind(graph.to_string())
|
||||
.bind(visibility)
|
||||
.bind(public_id)
|
||||
.bind(if allow_admin { 1_i64 } else { 0_i64 })
|
||||
.bind(user_id)
|
||||
.execute(&*state.db)
|
||||
.await?;
|
||||
if result.rows_affected() == 0 {
|
||||
return Err(AppError::NotFound("Workflow not found".into()));
|
||||
}
|
||||
get_workflow(&state.db, public_id, user_id, allow_admin).await
|
||||
}
|
||||
|
||||
pub async fn delete_workflow(
|
||||
state: &AppState,
|
||||
public_id: &str,
|
||||
user_id: Option<i64>,
|
||||
allow_admin: bool,
|
||||
) -> AppResult<()> {
|
||||
let result = sqlx::query(
|
||||
"UPDATE workflows
|
||||
SET deleted_at = datetime('now'), updated_at = datetime('now')
|
||||
WHERE public_id = ?
|
||||
AND deleted_at IS NULL
|
||||
AND (? = 1 OR user_id = ?)",
|
||||
)
|
||||
.bind(public_id)
|
||||
.bind(if allow_admin { 1_i64 } else { 0_i64 })
|
||||
.bind(user_id)
|
||||
.execute(&*state.db)
|
||||
.await?;
|
||||
if result.rows_affected() == 0 {
|
||||
return Err(AppError::NotFound("Workflow not found".into()));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn insert_run(
|
||||
state: &AppState,
|
||||
workflow_id: i64,
|
||||
user_id: Option<i64>,
|
||||
input_overrides: &Value,
|
||||
) -> AppResult<(i64, String)> {
|
||||
let public_id = new_run_id();
|
||||
let result = sqlx::query(
|
||||
"INSERT INTO workflow_runs
|
||||
(public_id, workflow_id, user_id, status, input_overrides, started_at)
|
||||
VALUES (?, ?, ?, 'running', ?, datetime('now'))",
|
||||
)
|
||||
.bind(&public_id)
|
||||
.bind(workflow_id)
|
||||
.bind(user_id)
|
||||
.bind(input_overrides.to_string())
|
||||
.execute(&*state.db)
|
||||
.await?;
|
||||
let id = result.last_insert_rowid();
|
||||
Ok((id, public_id))
|
||||
}
|
||||
|
||||
pub async fn insert_node_run_started(
|
||||
state: &AppState,
|
||||
run_id: i64,
|
||||
node_id: &str,
|
||||
node_type: &str,
|
||||
inputs: &Value,
|
||||
) -> AppResult<i64> {
|
||||
let result = sqlx::query(
|
||||
"INSERT INTO workflow_node_runs
|
||||
(run_id, node_id, node_type, status, inputs, started_at)
|
||||
VALUES (?, ?, ?, 'running', ?, datetime('now'))",
|
||||
)
|
||||
.bind(run_id)
|
||||
.bind(node_id)
|
||||
.bind(node_type)
|
||||
.bind(inputs.to_string())
|
||||
.execute(&*state.db)
|
||||
.await?;
|
||||
let id = result.last_insert_rowid();
|
||||
Ok(id)
|
||||
}
|
||||
|
||||
pub async fn complete_node_run(
|
||||
state: &AppState,
|
||||
node_run_id: i64,
|
||||
outputs: &Value,
|
||||
) -> AppResult<()> {
|
||||
sqlx::query(
|
||||
"UPDATE workflow_node_runs
|
||||
SET status = 'completed', outputs = ?, completed_at = datetime('now')
|
||||
WHERE id = ?",
|
||||
)
|
||||
.bind(outputs.to_string())
|
||||
.bind(node_run_id)
|
||||
.execute(&*state.db)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn fail_node_run(state: &AppState, node_run_id: i64, error: &str) -> AppResult<()> {
|
||||
sqlx::query(
|
||||
"UPDATE workflow_node_runs
|
||||
SET status = 'failed', error_message = ?, completed_at = datetime('now')
|
||||
WHERE id = ?",
|
||||
)
|
||||
.bind(error)
|
||||
.bind(node_run_id)
|
||||
.execute(&*state.db)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn complete_run(state: &AppState, run_id: i64, outputs: &Value) -> AppResult<()> {
|
||||
sqlx::query(
|
||||
"UPDATE workflow_runs
|
||||
SET status = 'completed', final_outputs = ?, completed_at = datetime('now')
|
||||
WHERE id = ?",
|
||||
)
|
||||
.bind(outputs.to_string())
|
||||
.bind(run_id)
|
||||
.execute(&*state.db)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn fail_run(state: &AppState, run_id: i64, error: &str) -> AppResult<()> {
|
||||
sqlx::query(
|
||||
"UPDATE workflow_runs
|
||||
SET status = 'failed', error_message = ?, completed_at = datetime('now')
|
||||
WHERE id = ?",
|
||||
)
|
||||
.bind(error)
|
||||
.bind(run_id)
|
||||
.execute(&*state.db)
|
||||
.await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn delete_runs(
|
||||
state: &AppState,
|
||||
public_ids: &[String],
|
||||
user_id: Option<i64>,
|
||||
allow_admin: bool,
|
||||
) -> AppResult<usize> {
|
||||
let mut deleted = 0_usize;
|
||||
for public_id in public_ids {
|
||||
sqlx::query(
|
||||
"DELETE FROM workflow_node_runs
|
||||
WHERE run_id IN (
|
||||
SELECT id
|
||||
FROM workflow_runs
|
||||
WHERE public_id = ?
|
||||
AND (? = 1 OR user_id = ?)
|
||||
)",
|
||||
)
|
||||
.bind(public_id)
|
||||
.bind(if allow_admin { 1_i64 } else { 0_i64 })
|
||||
.bind(user_id)
|
||||
.execute(&*state.db)
|
||||
.await?;
|
||||
let result = sqlx::query(
|
||||
"DELETE FROM workflow_runs
|
||||
WHERE public_id = ?
|
||||
AND (? = 1 OR user_id = ?)",
|
||||
)
|
||||
.bind(public_id)
|
||||
.bind(if allow_admin { 1_i64 } else { 0_i64 })
|
||||
.bind(user_id)
|
||||
.execute(&*state.db)
|
||||
.await?;
|
||||
deleted += result.rows_affected() as usize;
|
||||
}
|
||||
Ok(deleted)
|
||||
}
|
||||
|
||||
pub async fn clear_runs(
|
||||
state: &AppState,
|
||||
user_id: Option<i64>,
|
||||
allow_admin: bool,
|
||||
) -> AppResult<u64> {
|
||||
sqlx::query(
|
||||
"DELETE FROM workflow_node_runs
|
||||
WHERE run_id IN (
|
||||
SELECT id
|
||||
FROM workflow_runs
|
||||
WHERE (? = 1 OR user_id = ?)
|
||||
)",
|
||||
)
|
||||
.bind(if allow_admin { 1_i64 } else { 0_i64 })
|
||||
.bind(user_id)
|
||||
.execute(&*state.db)
|
||||
.await?;
|
||||
let result = sqlx::query(
|
||||
"DELETE FROM workflow_runs
|
||||
WHERE (? = 1 OR user_id = ?)",
|
||||
)
|
||||
.bind(if allow_admin { 1_i64 } else { 0_i64 })
|
||||
.bind(user_id)
|
||||
.execute(&*state.db)
|
||||
.await?;
|
||||
Ok(result.rows_affected())
|
||||
}
|
||||
+39
-2
@@ -18,11 +18,48 @@ export default defineConfig(({ mode }) => {
|
||||
build: {
|
||||
outDir: '../../dist/web',
|
||||
emptyOutDir: true,
|
||||
chunkSizeWarningLimit: 1100,
|
||||
rollupOptions: {
|
||||
output: {
|
||||
manualChunks(id) {
|
||||
if (id.includes('@visactor/react-vchart') || id.includes('/@visactor/')) {
|
||||
return 'vchart-vendor';
|
||||
if (!id.includes('/node_modules/')) {
|
||||
return undefined;
|
||||
}
|
||||
if (id.includes('/react/') || id.includes('/react-dom/') || id.includes('/scheduler/')) {
|
||||
return 'react-vendor';
|
||||
}
|
||||
if (id.includes('/react-router') || id.includes('/react-router-dom/')) {
|
||||
return 'router-vendor';
|
||||
}
|
||||
if (id.includes('/@dnd-kit/')) {
|
||||
return 'dnd-vendor';
|
||||
}
|
||||
if (id.includes('/mermaid/')) {
|
||||
return 'mermaid-vendor';
|
||||
}
|
||||
if (id.includes('/marked/')) {
|
||||
return 'markdown-vendor';
|
||||
}
|
||||
if (id.includes('/better-sqlite3/') || id.includes('/pg/') || id.includes('/mysql2/')) {
|
||||
return 'db-vendor';
|
||||
}
|
||||
if (id.includes('/@visactor/react-vchart/')) {
|
||||
return 'vchart-react-vendor';
|
||||
}
|
||||
if (id.includes('/@visactor/vchart/')) {
|
||||
return 'vchart-core-vendor';
|
||||
}
|
||||
if (id.includes('/@visactor/vgrammar')) {
|
||||
return 'vchart-grammar-vendor';
|
||||
}
|
||||
if (id.includes('/@visactor/vrender')) {
|
||||
return 'vchart-render-vendor';
|
||||
}
|
||||
if (id.includes('/@visactor/vdataset/') || id.includes('/@visactor/vutils/') || id.includes('/@visactor/vscale/') || id.includes('/@visactor/vutils-extension/')) {
|
||||
return 'vchart-utils-vendor';
|
||||
}
|
||||
if (id.includes('/@visactor/')) {
|
||||
return 'vchart-extra-vendor';
|
||||
}
|
||||
return undefined;
|
||||
},
|
||||
|
||||
Reference in New Issue
Block a user