feat: merge opus comfyui and login updates
This commit is contained in:
boos
@@ -0,0 +1,23 @@
|
|||||||
|
# Cargo build acceleration.
|
||||||
|
#
|
||||||
|
# The biggest single speedup is the `mold` linker (~3-5x faster link
|
||||||
|
# time on debug builds). It isn't installed by default on this host
|
||||||
|
# (no apt-get, see build-and-run.sh notes), so the mold block below is
|
||||||
|
# commented out — uncomment it after running:
|
||||||
|
#
|
||||||
|
# sudo apt-get install -y mold
|
||||||
|
#
|
||||||
|
# (or build it from source: https://github.com/rui314/mold)
|
||||||
|
#
|
||||||
|
# Without mold, the [profile.dev] tweaks in Cargo.toml still give a
|
||||||
|
# meaningful speedup (line-tables-only debuginfo + bumped codegen-units).
|
||||||
|
|
||||||
|
# Uncomment once mold is on PATH:
|
||||||
|
# [target.x86_64-unknown-linux-gnu]
|
||||||
|
# linker = "clang"
|
||||||
|
# rustflags = ["-C", "link-arg=-fuse-ld=mold"]
|
||||||
|
|
||||||
|
[build]
|
||||||
|
# Default to all CPUs for parallel codegen. Override with --jobs N if
|
||||||
|
# you need to leave headroom for other work.
|
||||||
|
# jobs = 8
|
||||||
@@ -17,4 +17,6 @@ docker/data/
|
|||||||
docker/.env
|
docker/.env
|
||||||
.ace-tool/
|
.ace-tool/
|
||||||
.comfyui/
|
.comfyui/
|
||||||
|
.claude/
|
||||||
|
.codex/
|
||||||
boosapi
|
boosapi
|
||||||
|
|||||||
@@ -0,0 +1,4 @@
|
|||||||
|
[submodule "comfy_source"]
|
||||||
|
path = comfy_source
|
||||||
|
url = https://github.com/comfyanonymous/ComfyUI.git
|
||||||
|
branch = master
|
||||||
@@ -31,9 +31,15 @@ boosapi-rs/ ← 根目录（已迁移至此）
|
|||||||
|
|
||||||
## ComfyUI Agent
|
## ComfyUI Agent
|
||||||
|
|
||||||
ComfyUI Agent 后端直连上游 API（`api.boos.lat/v1`），不走本地代理路由。
|
ComfyUI Agent 默认走 **本地模式**：复用本服务自身的 `/v1/*` 路由，bearer 用启动时引导出来的内置 downstream key（`name='comfyui-agent' AND user_id IS NULL`，密钥格式 `sk-comfyui-{uuid_v4_hex}`，同时镜像写入 `settings.builtin_key.comfyui_agent.secret`）。
|
||||||
|
|
||||||
支持的结构化标记：
|
行为契约：
|
||||||
|
- **Local（默认）**：`upstream_base` = `http://{server.host}/v1`（host 是 `0.0.0.0` / `::` 时改写为 `127.0.0.1`），bearer = 内置 key；`/api/comfyui-agent/status` 回 `mode: "local"` + `keyName: "comfyui-agent"`
|
||||||
|
- **FallbackRemote**：当 `BOOSAPI_UPSTREAM_BASE`（trim 后非空）被设置时切换到外部上游；`BOOSAPI_UPSTREAM_KEY` 缺省时仍用内置 key；status 回 `mode: "fallback_remote"` + `keyName: "external"`
|
||||||
|
- 内置 key 受保护：admin 列表对其 key 字段打码（`••••<后4位>`），删除返 403，禁用（`enabled=false`）返 403；`builtin_key.*` settings 不在 list 中暴露，PUT 返 403
|
||||||
|
- 启动期 `ensure_comfyui_agent_key` 是幂等的：缺密钥就生成、`enabled=0` 自动恢复并 `tracing::warn!`、settings 与 downstream_api_keys 双向同步
|
||||||
|
|
||||||
|
支持的结构化标记（输出由 LLM 生成、由 agent SSE 解析下发）：
|
||||||
- 【CHARACTER】 → 自动调用 gpt-image-2 生成四视图
|
- 【CHARACTER】 → 自动调用 gpt-image-2 生成四视图
|
||||||
- 【SCENE】 → 自动生成场景图
|
- 【SCENE】 → 自动生成场景图
|
||||||
- 【TTS】 → 自动调用 TTS API
|
- 【TTS】 → 自动调用 TTS API
|
||||||
@@ -58,3 +64,9 @@ BOOSAPI_AUTH__AUTH_TOKEN=change-me-admin-token
|
|||||||
前端为 React SPA，使用 Vite 构建，产物在 `dist/web/`：
|
前端为 React SPA，使用 Vite 构建，产物在 `dist/web/`：
|
||||||
- 开发时用 Vite dev server（需 Node.js）
|
- 开发时用 Vite dev server（需 Node.js）
|
||||||
- 生产时由 Rust 后端直接 serve 静态文件
|
- 生产时由 Rust 后端直接 serve 静态文件
|
||||||
|
|
||||||
|
### 登录与会话
|
||||||
|
|
||||||
|
- 登录页位于 `src/web/pages/auth/`：`LoginShell` 提供共享外壳，`AdminLoginPage`（token / 账密双模式）与 `UserLoginPage`（账密）分别承担两类受众。OAuth 入口用 `OAuthSlot` 渲染占位区域，**实际行为由后续 change 接入**——当前不要在这里挂任何 onClick / 跳转。
|
||||||
|
- 会话状态统一走 `src/web/authSession.ts` 的 `commitSession / clearSession / getSessionToken`，受众通过 `'admin' | 'user'` 区分；旧的 `persistAuthSession / clearAuthSession / getAuthToken / hasValidAuthSession`（及 `*UserSession` 同族）保留为 `@deprecated` 薄壳，新代码不要再引用。
|
||||||
|
- 登录错误文案统一走 `src/web/loginError.ts` 的 `resolveLoginErrorMessage(status, raw)`；新增登录错误请加在该处而不是各自分支里手写。
|
||||||
|
|||||||
+30
-2
@@ -4,9 +4,18 @@ version = "1.3.0"
|
|||||||
edition = "2021"
|
edition = "2021"
|
||||||
description = "AI API Gateway — Rust rewrite of BoosAPI"
|
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]
|
[dependencies]
|
||||||
# Web framework
|
# Web framework
|
||||||
axum = { version = "0.8", features = ["macros"] }
|
axum = { version = "0.8", features = ["macros", "ws", "multipart"] }
|
||||||
tower = "0.5"
|
tower = "0.5"
|
||||||
tower-http = { version = "0.6", features = ["cors", "trace", "fs", "compression-gzip"] }
|
tower-http = { version = "0.6", features = ["cors", "trace", "fs", "compression-gzip"] }
|
||||||
|
|
||||||
@@ -23,7 +32,7 @@ sqlx = { version = "0.8", features = [
|
|||||||
] }
|
] }
|
||||||
|
|
||||||
# HTTP client
|
# HTTP client
|
||||||
reqwest = { version = "0.12", features = ["stream", "json"] }
|
reqwest = { version = "0.12", features = ["stream", "json", "multipart"] }
|
||||||
lettre = { version = "0.11", default-features = false, features = ["builder", "hostname", "pool", "rustls-tls", "smtp-transport"] }
|
lettre = { version = "0.11", default-features = false, features = ["builder", "hostname", "pool", "rustls-tls", "smtp-transport"] }
|
||||||
|
|
||||||
# Serialization
|
# Serialization
|
||||||
@@ -66,6 +75,11 @@ once_cell = "1"
|
|||||||
parking_lot = "0.12"
|
parking_lot = "0.12"
|
||||||
mime = "0.3"
|
mime = "0.3"
|
||||||
|
|
||||||
|
# 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"
|
||||||
|
|
||||||
[dev-dependencies]
|
[dev-dependencies]
|
||||||
rstest = "0.24"
|
rstest = "0.24"
|
||||||
wiremock = "0.6"
|
wiremock = "0.6"
|
||||||
@@ -77,3 +91,17 @@ axum-test = "16"
|
|||||||
opt-level = 3
|
opt-level = 3
|
||||||
lto = true
|
lto = true
|
||||||
codegen-units = 1
|
codegen-units = 1
|
||||||
|
|
||||||
|
# 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
|
||||||
|
# build with no debuginfo at all since we rarely step into them.
|
||||||
|
[profile.dev]
|
||||||
|
opt-level = 0
|
||||||
|
debug = "line-tables-only"
|
||||||
|
codegen-units = 256
|
||||||
|
incremental = true
|
||||||
|
|
||||||
|
[profile.dev.package."*"]
|
||||||
|
opt-level = 0
|
||||||
|
debug = false
|
||||||
|
|||||||
@@ -70,7 +70,8 @@
|
|||||||
- **ComfyUI 工作流导出**：完整 workflow JSON，可直接导入 ComfyUI
|
- **ComfyUI 工作流导出**：完整 workflow JSON，可直接导入 ComfyUI
|
||||||
- **Rust 后端**（boosapi-rs）：高性能 SSE 流式响应，零 Node.js 运行时依赖
|
- **Rust 后端**（boosapi-rs）：高性能 SSE 流式响应，零 Node.js 运行时依赖
|
||||||
- 侧边栏「视频助手」入口，与聚合网关同一界面
|
- 侧边栏「视频助手」入口，与聚合网关同一界面
|
||||||
- 自动对接上游 API（`api.boos.lat/v1`），无需额外配置
|
- **默认走本地模式**：复用本服务自身的 `/v1/*` 路由 + 启动时引导出来的内置 downstream key（`comfyui-agent`），开箱即用，无需额外配置
|
||||||
|
- **回退远端**：设置 `BOOSAPI_UPSTREAM_BASE`（可选 `BOOSAPI_UPSTREAM_KEY`）即切换为外部上游；`/api/comfyui-agent/status` 的 `mode` 字段会反映当前模式
|
||||||
|
|
||||||
### 🌐 统一代理网关
|
### 🌐 统一代理网关
|
||||||
|
|
||||||
@@ -228,7 +229,10 @@ docker compose up -d
|
|||||||
│ └── notification/ # 通知推送
|
│ └── notification/ # 通知推送
|
||||||
├── migrations/ # SQLite 迁移
|
├── migrations/ # SQLite 迁移
|
||||||
├── src/web/ # 前端源码（React 19 + Vite）
|
├── src/web/ # 前端源码（React 19 + Vite）
|
||||||
|
│ ├── authSession.ts # 统一会话工具：commitSession/clearSession/getSessionToken（按 'admin' | 'user' 分桶）
|
||||||
|
│ ├── loginError.ts # 登录错误归一化（resolveLoginErrorMessage）
|
||||||
│ ├── pages/
|
│ ├── pages/
|
||||||
|
│ │ ├── auth/ # 登录页拆分：LoginShell + AdminLoginPage + UserLoginPage（OAuthSlot 占位，行为由后续 change 接入）
|
||||||
│ │ ├── ComfyUIAgent.tsx # 视频助手聊天界面
|
│ │ ├── ComfyUIAgent.tsx # 视频助手聊天界面
|
||||||
│ │ └── ...
|
│ │ └── ...
|
||||||
│ └── components/
|
│ └── components/
|
||||||
@@ -253,7 +257,8 @@ BOOSAPI_AUTH__JWT_SECRET=your-jwt-secret
|
|||||||
# 数据库
|
# 数据库
|
||||||
BOOSAPI_DB__URL=sqlite:data/boosapi.db?mode=rwc
|
BOOSAPI_DB__URL=sqlite:data/boosapi.db?mode=rwc
|
||||||
|
|
||||||
# 上游 API（ComfyUI Agent 使用）
|
# 上游 API（ComfyUI Agent 回退用，默认本地模式时不需要设置）
|
||||||
|
# 设置 BOOSAPI_UPSTREAM_BASE 后切换到 fallback-remote 模式；BOOSAPI_UPSTREAM_KEY 缺省时仍用内置 comfyui-agent key
|
||||||
BOOSAPI_UPSTREAM_BASE=https://api.boos.lat/v1
|
BOOSAPI_UPSTREAM_BASE=https://api.boos.lat/v1
|
||||||
BOOSAPI_UPSTREAM_KEY=sk-your-key
|
BOOSAPI_UPSTREAM_KEY=sk-your-key
|
||||||
|
|
||||||
|
|||||||
@@ -1,298 +0,0 @@
|
|||||||
"""
|
|
||||||
BoosAPI Custom Nodes for ComfyUI
|
|
||||||
==================================
|
|
||||||
Call local BoosAPI endpoints directly from ComfyUI workflows:
|
|
||||||
- Text generation (gpt-5.5)
|
|
||||||
- Image generation (gpt-image-2)
|
|
||||||
- Text-to-Speech
|
|
||||||
- Video generation (async)
|
|
||||||
|
|
||||||
Default: http://127.0.0.1:4000/v1 (same container)
|
|
||||||
Override via env:
|
|
||||||
BOOSAPI_BASE = https://api.boos.lat/v1 (standalone ComfyUI)
|
|
||||||
BOOSAPI_KEY=***
|
|
||||||
"""
|
|
||||||
|
|
||||||
import os
|
|
||||||
import json
|
|
||||||
import time
|
|
||||||
import requests
|
|
||||||
import torch
|
|
||||||
import numpy as np
|
|
||||||
from PIL import Image
|
|
||||||
from io import BytesIO
|
|
||||||
import folder_paths
|
|
||||||
|
|
||||||
# ── Config ──────────────────────────────────────────────────────────────────
|
|
||||||
|
|
||||||
# When running inside the same container as BoosAPI, use localhost.
|
|
||||||
# For standalone ComfyUI installs, set BOOSAPI_BASE env to https://api.boos.lat/v1
|
|
||||||
BOOSAPI_BASE = os.getenv("BOOSAPI_BASE", "https://api.boos.lat/v1")
|
|
||||||
BOOSAPI_KEY = os.getenv("BOOSAPI_KEY", "sk-boos4721")
|
|
||||||
|
|
||||||
HEADERS = {
|
|
||||||
"Authorization": f"Bearer {BOOSAPI_KEY}",
|
|
||||||
"Content-Type": "application/json",
|
|
||||||
}
|
|
||||||
|
|
||||||
|
|
||||||
def _call_api(endpoint: str, payload: dict, timeout: int = 120) -> dict:
|
|
||||||
url = f"{BOOSAPI_BASE}/{endpoint.lstrip('/')}"
|
|
||||||
resp = requests.post(url, headers=HEADERS, json=payload, timeout=timeout)
|
|
||||||
if not resp.ok:
|
|
||||||
raise RuntimeError(f"BoosAPI {endpoint} failed HTTP {resp.status_code}: {resp.text}")
|
|
||||||
return resp.json()
|
|
||||||
|
|
||||||
|
|
||||||
# ── Shared helpers ──────────────────────────────────────────────────────────
|
|
||||||
|
|
||||||
def _pil_to_tensor(img: Image.Image) -> torch.Tensor:
|
|
||||||
"""PIL -> ComfyUI IMAGE (1, H, W, 3) float32 0-1."""
|
|
||||||
img = img.convert("RGB")
|
|
||||||
arr = np.array(img).astype(np.float32) / 255.0
|
|
||||||
return torch.from_numpy(arr).unsqueeze(0)
|
|
||||||
|
|
||||||
|
|
||||||
def _tensor_to_pil(tensor: torch.Tensor) -> Image.Image:
|
|
||||||
"""ComfyUI IMAGE -> PIL."""
|
|
||||||
arr = (tensor.squeeze(0).cpu().numpy() * 255).astype(np.uint8)
|
|
||||||
return Image.fromarray(arr)
|
|
||||||
|
|
||||||
|
|
||||||
# ── Node: Text Generation (gpt-5.5) ─────────────────────────────────────────
|
|
||||||
|
|
||||||
class BoosAPITextGen:
|
|
||||||
"""Call gpt-5.5 chat completion. Useful for prompt engineering within workflows."""
|
|
||||||
|
|
||||||
@classmethod
|
|
||||||
def INPUT_TYPES(cls):
|
|
||||||
return {
|
|
||||||
"required": {
|
|
||||||
"prompt": ("STRING", {"multiline": True, "default": ""}),
|
|
||||||
"model": ("STRING", {"default": "gpt-5.5"}),
|
|
||||||
"temperature": ("FLOAT", {"default": 0.7, "min": 0, "max": 2, "step": 0.1}),
|
|
||||||
"max_tokens": ("INT", {"default": 2048, "min": 1, "max": 32768}),
|
|
||||||
},
|
|
||||||
"optional": {
|
|
||||||
"system_prompt": ("STRING", {"multiline": True, "default": ""}),
|
|
||||||
},
|
|
||||||
}
|
|
||||||
|
|
||||||
RETURN_TYPES = ("STRING",)
|
|
||||||
RETURN_NAMES = ("text",)
|
|
||||||
FUNCTION = "generate"
|
|
||||||
CATEGORY = "BoosAPI"
|
|
||||||
|
|
||||||
def generate(self, prompt, model, temperature, max_tokens, system_prompt=""):
|
|
||||||
messages = []
|
|
||||||
if system_prompt:
|
|
||||||
messages.append({"role": "system", "content": system_prompt})
|
|
||||||
messages.append({"role": "user", "content": prompt})
|
|
||||||
|
|
||||||
data = _call_api("chat/completions", {
|
|
||||||
"model": model,
|
|
||||||
"messages": messages,
|
|
||||||
"temperature": temperature,
|
|
||||||
"max_tokens": max_tokens,
|
|
||||||
})
|
|
||||||
text = data["choices"][0]["message"]["content"]
|
|
||||||
return (text,)
|
|
||||||
|
|
||||||
|
|
||||||
# ── Node: Image Generation (gpt-image-2) ────────────────────────────────────
|
|
||||||
|
|
||||||
class BoosAPIImageGen:
|
|
||||||
"""Generate image via BoosAPI. Replaces CheckpointLoader+KSampler+VAE chain."""
|
|
||||||
|
|
||||||
@classmethod
|
|
||||||
def INPUT_TYPES(cls):
|
|
||||||
return {
|
|
||||||
"required": {
|
|
||||||
"prompt": ("STRING", {"multiline": True, "default": ""}),
|
|
||||||
"model": (["gpt-image-2", "dall-e-3"], {"default": "gpt-image-2"}),
|
|
||||||
"size": (["1024x1024", "1792x1024", "1024x1792"], {"default": "1024x1024"}),
|
|
||||||
},
|
|
||||||
"optional": {
|
|
||||||
"negative_prompt": ("STRING", {"multiline": True, "default": ""}),
|
|
||||||
"n": ("INT", {"default": 1, "min": 1, "max": 4}),
|
|
||||||
},
|
|
||||||
}
|
|
||||||
|
|
||||||
RETURN_TYPES = ("IMAGE", "STRING")
|
|
||||||
RETURN_NAMES = ("image", "image_url")
|
|
||||||
FUNCTION = "generate"
|
|
||||||
CATEGORY = "BoosAPI"
|
|
||||||
|
|
||||||
def generate(self, prompt, model, size, negative_prompt="", n=1):
|
|
||||||
payload = {
|
|
||||||
"model": model,
|
|
||||||
"prompt": prompt,
|
|
||||||
"n": n,
|
|
||||||
"size": size,
|
|
||||||
}
|
|
||||||
if negative_prompt:
|
|
||||||
payload["negative_prompt"] = negative_prompt
|
|
||||||
|
|
||||||
data = _call_api("images/generations", payload)
|
|
||||||
|
|
||||||
# Download first image
|
|
||||||
image_url = data["data"][0]["url"]
|
|
||||||
img_resp = requests.get(image_url, timeout=60)
|
|
||||||
img = Image.open(BytesIO(img_resp.content))
|
|
||||||
tensor = _pil_to_tensor(img)
|
|
||||||
return (tensor, image_url)
|
|
||||||
|
|
||||||
|
|
||||||
# ── Node: Text-to-Speech ────────────────────────────────────────────────────
|
|
||||||
|
|
||||||
class BoosAPITTS:
|
|
||||||
"""Generate speech audio from text via BoosAPI TTS."""
|
|
||||||
|
|
||||||
@classmethod
|
|
||||||
def INPUT_TYPES(cls):
|
|
||||||
return {
|
|
||||||
"required": {
|
|
||||||
"text": ("STRING", {"multiline": True, "default": ""}),
|
|
||||||
"voice": (["alloy", "echo", "fable", "onyx", "nova", "shimmer"], {"default": "alloy"}),
|
|
||||||
},
|
|
||||||
"optional": {
|
|
||||||
"speed": ("FLOAT", {"default": 1.0, "min": 0.25, "max": 4.0, "step": 0.25}),
|
|
||||||
},
|
|
||||||
}
|
|
||||||
|
|
||||||
RETURN_TYPES = ("STRING",) # file path to generated audio
|
|
||||||
RETURN_NAMES = ("audio_path",)
|
|
||||||
FUNCTION = "generate"
|
|
||||||
CATEGORY = "BoosAPI"
|
|
||||||
OUTPUT_NODE = True
|
|
||||||
|
|
||||||
def generate(self, text, voice, speed=1.0):
|
|
||||||
url = f"{BOOSAPI_BASE}/audio/speech"
|
|
||||||
resp = requests.post(
|
|
||||||
url, headers=HEADERS,
|
|
||||||
json={"model": "tts-1", "input": text, "voice": voice, "speed": speed},
|
|
||||||
timeout=120,
|
|
||||||
)
|
|
||||||
if not resp.ok:
|
|
||||||
raise RuntimeError(f"TTS failed HTTP {resp.status_code}: {resp.text}")
|
|
||||||
|
|
||||||
# Save to ComfyUI output directory
|
|
||||||
output_dir = folder_paths.get_output_directory()
|
|
||||||
os.makedirs(output_dir, exist_ok=True)
|
|
||||||
filename = f"boosapi_tts_{int(time.time())}.mp3"
|
|
||||||
filepath = os.path.join(output_dir, filename)
|
|
||||||
with open(filepath, "wb") as f:
|
|
||||||
f.write(resp.content)
|
|
||||||
|
|
||||||
print(f"[BoosAPI] TTS saved: {filepath}")
|
|
||||||
return (filepath,)
|
|
||||||
|
|
||||||
|
|
||||||
# ── Node: Video Generation (async) ─────────────────────────────────────────
|
|
||||||
|
|
||||||
class BoosAPIVideoGen:
|
|
||||||
"""Generate video via BoosAPI (async). Returns task ID for polling."""
|
|
||||||
|
|
||||||
@classmethod
|
|
||||||
def INPUT_TYPES(cls):
|
|
||||||
return {
|
|
||||||
"required": {
|
|
||||||
"prompt": ("STRING", {"multiline": True, "default": ""}),
|
|
||||||
"model": ("STRING", {"default": "cogvideo-5b"}),
|
|
||||||
"size": (["1024x1024", "1024x576", "576x1024", "1920x1080"], {"default": "1024x576"}),
|
|
||||||
},
|
|
||||||
"optional": {
|
|
||||||
"negative_prompt": ("STRING", {"multiline": True, "default": ""}),
|
|
||||||
"duration": ("INT", {"default": 5, "min": 2, "max": 30}),
|
|
||||||
"fps": ("INT", {"default": 16, "min": 8, "max": 30}),
|
|
||||||
},
|
|
||||||
}
|
|
||||||
|
|
||||||
RETURN_TYPES = ("STRING", "STRING")
|
|
||||||
RETURN_NAMES = ("task_id", "status_url")
|
|
||||||
FUNCTION = "generate"
|
|
||||||
CATEGORY = "BoosAPI"
|
|
||||||
|
|
||||||
def generate(self, prompt, model, size, negative_prompt="", duration=5, fps=16):
|
|
||||||
width, height = size.split("x")
|
|
||||||
payload = {
|
|
||||||
"model": model,
|
|
||||||
"prompt": prompt,
|
|
||||||
"width": int(width),
|
|
||||||
"height": int(height),
|
|
||||||
"duration": duration,
|
|
||||||
"fps": fps,
|
|
||||||
}
|
|
||||||
if negative_prompt:
|
|
||||||
payload["negative_prompt"] = negative_prompt
|
|
||||||
|
|
||||||
data = _call_api("videos", payload)
|
|
||||||
task_id = data.get("id", "")
|
|
||||||
status_url = f"{BOOSAPI_BASE}/videos/{task_id}"
|
|
||||||
return (task_id, status_url)
|
|
||||||
|
|
||||||
|
|
||||||
class BoosAPIVideoStatus:
|
|
||||||
"""Poll video generation status. Returns download URL when done."""
|
|
||||||
|
|
||||||
@classmethod
|
|
||||||
def INPUT_TYPES(cls):
|
|
||||||
return {
|
|
||||||
"required": {
|
|
||||||
"task_id": ("STRING", {"default": ""}),
|
|
||||||
},
|
|
||||||
"optional": {
|
|
||||||
"wait": ("BOOLEAN", {"default": True, "label_on": "Wait until done", "label_off": "Check once"}),
|
|
||||||
"poll_interval": ("INT", {"default": 5, "min": 1, "max": 60}),
|
|
||||||
"max_polls": ("INT", {"default": 60, "min": 1, "max": 600}),
|
|
||||||
},
|
|
||||||
}
|
|
||||||
|
|
||||||
RETURN_TYPES = ("STRING", "STRING")
|
|
||||||
RETURN_NAMES = ("status", "video_url")
|
|
||||||
FUNCTION = "poll"
|
|
||||||
CATEGORY = "BoosAPI"
|
|
||||||
|
|
||||||
def poll(self, task_id, wait=True, poll_interval=5, max_polls=60):
|
|
||||||
url = f"{BOOSAPI_BASE}/videos/{task_id}"
|
|
||||||
|
|
||||||
for attempt in range(max_polls):
|
|
||||||
resp = requests.get(url, headers=HEADERS, timeout=30)
|
|
||||||
if not resp.ok:
|
|
||||||
return ("error", f"HTTP {resp.status_code}")
|
|
||||||
|
|
||||||
data = resp.json()
|
|
||||||
status = data.get("status", "unknown")
|
|
||||||
|
|
||||||
if status == "completed":
|
|
||||||
video_url = data.get("video_url") or data.get("output", {}).get("video_url", "")
|
|
||||||
return ("completed", video_url)
|
|
||||||
elif status in ("failed", "error"):
|
|
||||||
error_msg = data.get("error", str(data))
|
|
||||||
return (f"failed: {error_msg}", "")
|
|
||||||
|
|
||||||
if not wait:
|
|
||||||
return (status, "")
|
|
||||||
|
|
||||||
time.sleep(poll_interval)
|
|
||||||
|
|
||||||
return ("timeout", "")
|
|
||||||
|
|
||||||
|
|
||||||
# ── Registration ────────────────────────────────────────────────────────────
|
|
||||||
|
|
||||||
NODE_CLASS_MAPPINGS = {
|
|
||||||
"BoosAPITextGen": BoosAPITextGen,
|
|
||||||
"BoosAPIImageGen": BoosAPIImageGen,
|
|
||||||
"BoosAPITTS": BoosAPITTS,
|
|
||||||
"BoosAPIVideoGen": BoosAPIVideoGen,
|
|
||||||
"BoosAPIVideoStatus": BoosAPIVideoStatus,
|
|
||||||
}
|
|
||||||
|
|
||||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
||||||
"BoosAPITextGen": "BoosAPI Text Gen (gpt-5.5)",
|
|
||||||
"BoosAPIImageGen": "BoosAPI Image Gen (gpt-image-2)",
|
|
||||||
"BoosAPITTS": "BoosAPI Text-to-Speech",
|
|
||||||
"BoosAPIVideoGen": "BoosAPI Video Gen (async)",
|
|
||||||
"BoosAPIVideoStatus": "BoosAPI Video Status",
|
|
||||||
}
|
|
||||||
+84
-7
@@ -47,6 +47,88 @@ log() { echo "[$(date '+%H:%M:%S')] $*"; }
|
|||||||
ok() { echo "[$(date '+%H:%M:%S')] ✓ $*"; }
|
ok() { echo "[$(date '+%H:%M:%S')] ✓ $*"; }
|
||||||
fail() { echo "[$(date '+%H:%M:%S')] ✗ $*" >&2; exit 1; }
|
fail() { echo "[$(date '+%H:%M:%S')] ✗ $*" >&2; exit 1; }
|
||||||
|
|
||||||
|
server_port() {
|
||||||
|
if [ -n "${BOOSAPI_SERVER__PORT:-}" ]; then
|
||||||
|
printf '%s\n' "$BOOSAPI_SERVER__PORT"
|
||||||
|
return
|
||||||
|
fi
|
||||||
|
awk '
|
||||||
|
/^\[server\]/ { in_server = 1; next }
|
||||||
|
/^\[/ { in_server = 0 }
|
||||||
|
in_server && /^[[:space:]]*port[[:space:]]*=/ {
|
||||||
|
gsub(/[^0-9]/, "", $0)
|
||||||
|
print $0
|
||||||
|
exit
|
||||||
|
}
|
||||||
|
' "$ROOT_DIR/config.toml" 2>/dev/null || true
|
||||||
|
}
|
||||||
|
|
||||||
|
is_running_pid() {
|
||||||
|
local pid="$1"
|
||||||
|
[ -n "$pid" ] && kill -0 "$pid" 2>/dev/null
|
||||||
|
}
|
||||||
|
|
||||||
|
stop_pids() {
|
||||||
|
local label="$1"
|
||||||
|
shift || true
|
||||||
|
local pids=()
|
||||||
|
local pid
|
||||||
|
|
||||||
|
for pid in "$@"; do
|
||||||
|
if [ -n "$pid" ] && [ "$pid" != "$$" ] && is_running_pid "$pid"; then
|
||||||
|
pids+=("$pid")
|
||||||
|
fi
|
||||||
|
done
|
||||||
|
|
||||||
|
if [ "${#pids[@]}" -eq 0 ]; then
|
||||||
|
return 0
|
||||||
|
fi
|
||||||
|
|
||||||
|
log "停止旧进程 ${label}: ${pids[*]}"
|
||||||
|
kill -TERM "${pids[@]}" 2>/dev/null || true
|
||||||
|
|
||||||
|
for _ in 1 2 3 4 5; do
|
||||||
|
local alive=()
|
||||||
|
for pid in "${pids[@]}"; do
|
||||||
|
if is_running_pid "$pid"; then
|
||||||
|
alive+=("$pid")
|
||||||
|
fi
|
||||||
|
done
|
||||||
|
if [ "${#alive[@]}" -eq 0 ]; then
|
||||||
|
ok "旧进程已停止"
|
||||||
|
return 0
|
||||||
|
fi
|
||||||
|
sleep 1
|
||||||
|
pids=("${alive[@]}")
|
||||||
|
done
|
||||||
|
|
||||||
|
log "旧进程未及时退出，强制结束: ${pids[*]}"
|
||||||
|
kill -KILL "${pids[@]}" 2>/dev/null || true
|
||||||
|
}
|
||||||
|
|
||||||
|
stop_previous_server() {
|
||||||
|
local binary_pids=()
|
||||||
|
local port_pids=()
|
||||||
|
local port
|
||||||
|
|
||||||
|
while IFS= read -r pid; do
|
||||||
|
[ -n "$pid" ] && binary_pids+=("$pid")
|
||||||
|
done < <(pgrep -f "target/(debug|release)/boosapi" 2>/dev/null || true)
|
||||||
|
|
||||||
|
stop_pids "boosapi binary" "${binary_pids[@]}"
|
||||||
|
|
||||||
|
port="$(server_port)"
|
||||||
|
port="${port:-4000}"
|
||||||
|
if command -v fuser >/dev/null 2>&1; then
|
||||||
|
while IFS= read -r pid; do
|
||||||
|
[ -n "$pid" ] && port_pids+=("$pid")
|
||||||
|
done < <(fuser "${port}/tcp" 2>/dev/null | tr ' ' '\n' || true)
|
||||||
|
stop_pids "占用 ${port}/tcp" "${port_pids[@]}"
|
||||||
|
else
|
||||||
|
log "未找到 fuser，跳过端口 ${port}/tcp 占用检查"
|
||||||
|
fi
|
||||||
|
}
|
||||||
|
|
||||||
# ── 同步 ComfyUI 子模块 + 应用本地补丁 ─────────────────────────────────────────
|
# ── 同步 ComfyUI 子模块 + 应用本地补丁 ─────────────────────────────────────────
|
||||||
# comfy_source/ 是 vanilla ComfyUI 的 git 子模块，本地补丁存放在
|
# comfy_source/ 是 vanilla ComfyUI 的 git 子模块，本地补丁存放在
|
||||||
# comfy_patches/overlay/，由 comfy_patches/apply.sh 覆盖到子模块上。
|
# comfy_patches/overlay/，由 comfy_patches/apply.sh 覆盖到子模块上。
|
||||||
@@ -149,13 +231,8 @@ if [ "$NO_RUN" -eq 1 ] || [ "$FRONTEND_ONLY" -eq 1 ]; then
|
|||||||
exit 0
|
exit 0
|
||||||
fi
|
fi
|
||||||
|
|
||||||
# 停止旧进程
|
# 启动前清理旧后端进程和端口占用，避免 exec 后绑定端口失败。
|
||||||
OLD_PID=$(pgrep -f "target/(debug|release)/boosapi" || true)
|
stop_previous_server
|
||||||
if [ -n "$OLD_PID" ]; then
|
|
||||||
log "停止旧进程 (PID $OLD_PID)..."
|
|
||||||
kill "$OLD_PID" 2>/dev/null || true
|
|
||||||
sleep 1
|
|
||||||
fi
|
|
||||||
|
|
||||||
log "启动服务..."
|
log "启动服务..."
|
||||||
exec "$BINARY"
|
exec "$BINARY"
|
||||||
|
|||||||
@@ -0,0 +1,136 @@
|
|||||||
|
# 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).
|
||||||
Executable
+49
@@ -0,0 +1,49 @@
|
|||||||
|
#!/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/."
|
||||||
@@ -0,0 +1,479 @@
|
|||||||
|
<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)
|
||||||
@@ -0,0 +1,44 @@
|
|||||||
|
# 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.
|
||||||
@@ -0,0 +1,440 @@
|
|||||||
|
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
|
||||||
@@ -0,0 +1,98 @@
|
|||||||
|
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)
|
||||||
@@ -0,0 +1,841 @@
|
|||||||
|
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
@@ -0,0 +1,37 @@
|
|||||||
|
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
@@ -0,0 +1,289 @@
|
|||||||
|
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
|
||||||
@@ -0,0 +1,144 @@
|
|||||||
|
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())
|
||||||
|
)
|
||||||
Submodule
+1
Submodule comfy_source added at 26515acd23
@@ -0,0 +1,124 @@
|
|||||||
|
# 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"]
|
||||||
Executable
+34
@@ -0,0 +1,34 @@
|
|||||||
|
#!/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,16 +0,0 @@
|
|||||||
services:
|
|
||||||
boosapi:
|
|
||||||
image: 1467078763/boosapi:latest
|
|
||||||
ports:
|
|
||||||
- "127.0.0.1:4000:4000"
|
|
||||||
volumes:
|
|
||||||
- ./data:/app/data
|
|
||||||
environment:
|
|
||||||
AUTH_TOKEN: ${AUTH_TOKEN:?AUTH_TOKEN is required}
|
|
||||||
PROXY_TOKEN: ${PROXY_TOKEN:?PROXY_TOKEN is required}
|
|
||||||
CHECKIN_CRON: "0 8 * * *"
|
|
||||||
BALANCE_REFRESH_CRON: "0 * * * *"
|
|
||||||
PORT: ${PORT:-4000}
|
|
||||||
DATA_DIR: /app/data
|
|
||||||
TZ: ${TZ:-Asia/Shanghai}
|
|
||||||
restart: unless-stopped
|
|
||||||
@@ -0,0 +1,2 @@
|
|||||||
|
schema: spec-driven
|
||||||
|
created: 2026-05-20
|
||||||
@@ -0,0 +1,220 @@
|
|||||||
|
## Context
|
||||||
|
|
||||||
|
The frontend `/routes` page is end-to-end broken: `POST /api/routes/rebuild`
|
||||||
|
in `src/routes/admin/compat.rs:368-370` was a stub that returned
|
||||||
|
`{ success: true }` without touching `token_routes` / `route_channels`, and
|
||||||
|
`GET /api/models/token-candidates` was unregistered. Sites and accounts
|
||||||
|
configure cleanly, models populate `model_availability` via the existing
|
||||||
|
`POST /api/accounts/{id}/models/refresh` path, but nothing ever fanned that
|
||||||
|
out into route rows, so `TokenRoutes.tsx` rendered "暂无路由" no matter
|
||||||
|
what.
|
||||||
|
|
||||||
|
The proposal calls for three capabilities — `route-model-discovery`,
|
||||||
|
`route-rebuild-pipeline`, `route-candidate-api` — plus an auto-trigger
|
||||||
|
hook on account creation. This document records what was actually built
|
||||||
|
and the trade-offs we accepted relative to the proposal.
|
||||||
|
|
||||||
|
**Stakeholders:**
|
||||||
|
- Backend (this change): Rust handler authors maintaining `routes::admin`.
|
||||||
|
- Frontend (`TokenRoutes.tsx:411-412`, `:252`): consumes the rebuild
|
||||||
|
summary and the candidates payload — contracts are load-bearing.
|
||||||
|
- Operators: re-runs `POST /api/routes/rebuild` whenever model topology
|
||||||
|
shifts; relies on `manual_override = 1` channels surviving rebuilds.
|
||||||
|
|
||||||
|
**Constraints:**
|
||||||
|
- SQLite as the backing store; no schema migrations in this change.
|
||||||
|
- Existing `model_availability` table is the source of truth for model
|
||||||
|
presence — already filled by the unmodified
|
||||||
|
`POST /api/accounts/{id}/models/refresh` handler.
|
||||||
|
- `manual_override = 1` rows in `route_channels` are user-curated and
|
||||||
|
must survive every rebuild, untouched.
|
||||||
|
|
||||||
|
## Goals / Non-Goals
|
||||||
|
|
||||||
|
**Goals:**
|
||||||
|
- Turn `POST /api/routes/rebuild` into a real pipeline that reads
|
||||||
|
`model_availability` (filtered by `accounts.status` and
|
||||||
|
`site_disabled_models`) and writes `token_routes` + `route_channels`.
|
||||||
|
- Return a `RebuildSummary` whose camelCase fields match the
|
||||||
|
`TokenRoutes.tsx:411-412` toast contract (`createdRoutes`,
|
||||||
|
`createdChannels`).
|
||||||
|
- Register `GET /api/models/token-candidates` returning the four-field
|
||||||
|
shape `TokenRoutes.tsx:252` already calls.
|
||||||
|
- Fire a background rebuild after account creation so first-time users
|
||||||
|
don't have to click "重建" before the page works.
|
||||||
|
- Preserve `display_name` / `display_icon` / `routing_strategy` on
|
||||||
|
re-runs — manual edits on `token_routes` must not be clobbered.
|
||||||
|
|
||||||
|
**Non-Goals:**
|
||||||
|
- Live `/v1/models` discovery (`route-model-discovery` capability). The
|
||||||
|
proposal sketched per-platform fetchers (OpenAI / Anthropic / generic)
|
||||||
|
and a per-`(site, account)` cache. We did not build them. Reason: the
|
||||||
|
existing `model_availability` table is already filled by the refresh
|
||||||
|
endpoint, and "就用现成的" was the explicit call. Re-scoped as a
|
||||||
|
potential follow-up change.
|
||||||
|
- `modelsMissingTokenGroups` and `endpointTypesByModel` content. The
|
||||||
|
legacy Node handler computes these from a pricing-catalog fetcher and
|
||||||
|
a marketplace-cache file that don't exist in this codebase. Returning
|
||||||
|
`{}` keeps the contract intact; the frontend tolerates empty objects.
|
||||||
|
- New crates. We reused `sqlx`, `tokio`, and `axum` already in the tree.
|
||||||
|
|
||||||
|
## Decisions
|
||||||
|
|
||||||
|
### D1. Read from `model_availability` instead of fetching upstream
|
||||||
|
|
||||||
|
**Choice:** The rebuild reads `model_availability ⋈ accounts` (filtered
|
||||||
|
by `accounts.status = 'active'` and an `EXISTS NOT` against
|
||||||
|
`site_disabled_models`).
|
||||||
|
|
||||||
|
**Alternatives considered:**
|
||||||
|
- Build the per-platform fetcher trait the proposal described and call
|
||||||
|
it inline during rebuild. Rejected: doubles the scope, ads network
|
||||||
|
failure modes to a write-path that should be deterministic, and the
|
||||||
|
table is already kept current by the refresh endpoint.
|
||||||
|
- Fetch on-demand only when `model_availability` is empty for an
|
||||||
|
account. Rejected: hard to reason about — user can't tell whether a
|
||||||
|
rebuild used cached or fresh data.
|
||||||
|
|
||||||
|
**Rationale:** Separation of concerns. "Refresh models for an account"
|
||||||
|
is one job (already implemented); "rebuild routes from known models"
|
||||||
|
is another. Composing them via the table keeps each path testable.
|
||||||
|
|
||||||
|
### D2. Single transaction for the whole rebuild
|
||||||
|
|
||||||
|
**Choice:** `state.db.begin()` at the top, `tx.commit()` at the bottom.
|
||||||
|
Every read and write goes through `&mut *tx`.
|
||||||
|
|
||||||
|
**Alternatives:**
|
||||||
|
- Per-route transaction so partial failures still produce some routes.
|
||||||
|
Rejected: `token_routes` and `route_channels` get out of sync if we
|
||||||
|
abort halfway, which is exactly the bug we're fixing.
|
||||||
|
|
||||||
|
**Rationale:** SQLite write throughput is dominated by fsync, not lock
|
||||||
|
contention; one big transaction is faster than many small ones for the
|
||||||
|
sizes we expect (tens to low hundreds of accounts). Rebuild is rare
|
||||||
|
enough that holding the write lock for the duration is acceptable.
|
||||||
|
|
||||||
|
### D3. Reuse existing `token_routes` rows by `model_pattern`
|
||||||
|
|
||||||
|
**Choice:** Pre-load `(model_pattern → id)` into a `HashMap`, look up
|
||||||
|
each model name in the map, only `INSERT` when absent.
|
||||||
|
|
||||||
|
**Alternatives:**
|
||||||
|
- `INSERT ... ON CONFLICT DO NOTHING` and then `SELECT` the id back.
|
||||||
|
Rejected: requires a unique index on `model_pattern` we didn't want
|
||||||
|
to add as a side effect; also two round-trips per model.
|
||||||
|
- Truncate-and-rewrite. Rejected: drops user-edited
|
||||||
|
`display_name`/`display_icon`.
|
||||||
|
|
||||||
|
**Rationale:** O(1) lookup per model, one query upfront, manual edits
|
||||||
|
preserved.
|
||||||
|
|
||||||
|
### D4. `manual_override = 1` channels are sacred
|
||||||
|
|
||||||
|
**Choice:** When reconciling channels under a route, partition existing
|
||||||
|
rows into `manual_keys` (skipped entirely) and `auto_by_key`
|
||||||
|
(reconciled). The desired-set excludes any `(account, model)` covered
|
||||||
|
by a manual row, so we never insert an auto row that would shadow a
|
||||||
|
manual one.
|
||||||
|
|
||||||
|
**Rationale:** Operators set `manual_override = 1` precisely to opt out
|
||||||
|
of automation. Stomping it would erase intent.
|
||||||
|
|
||||||
|
### D5. `token_id` selection is best-effort, nullable-tolerant
|
||||||
|
|
||||||
|
**Choice:** When inserting a new `route_channels` row, pick the
|
||||||
|
account's default-flagged enabled token via
|
||||||
|
`ORDER BY COALESCE(is_default, 0) DESC, id ASC LIMIT 1`. If the account
|
||||||
|
has no tokens, store `NULL` and let the proxy layer's existing
|
||||||
|
"account-level token" fallback handle it at request time.
|
||||||
|
|
||||||
|
**Alternatives:**
|
||||||
|
- Skip channels for accounts with no tokens. Rejected: the proposal
|
||||||
|
surfaces these as `modelsWithoutToken` in the candidates API; they
|
||||||
|
should still appear in the route list so the operator sees them.
|
||||||
|
- Loop over every enabled token and create one channel per token.
|
||||||
|
Rejected: explodes `route_channels` and conflates "account routing"
|
||||||
|
with "token routing", which the schema treats as separate concerns.
|
||||||
|
|
||||||
|
### D6. Auto-rebuild via fire-and-forget `tokio::spawn`
|
||||||
|
|
||||||
|
**Choice:** `spawn_route_rebuild(state.clone())` after a successful
|
||||||
|
account insert in both the single-account and batch branches of
|
||||||
|
`POST /api/accounts`. Errors are logged at `warn`, the response is not
|
||||||
|
delayed.
|
||||||
|
|
||||||
|
**Alternatives:**
|
||||||
|
- Synchronous rebuild in the request path. Rejected: adds rebuild
|
||||||
|
latency to every account creation, including ones where models
|
||||||
|
haven't been refreshed yet (so the rebuild would be a no-op). Worst
|
||||||
|
case: a slow rebuild times out the HTTP request.
|
||||||
|
- Job queue / cron. Rejected: introduces infrastructure for a single
|
||||||
|
trigger; we already have the `/api/routes/rebuild` endpoint as the
|
||||||
|
manual fallback.
|
||||||
|
|
||||||
|
**Rationale:** The user's actual flow is "add account → refresh models
|
||||||
|
→ rebuild fires → routes appear". Auto-trigger handles step 3 without
|
||||||
|
blocking step 1.
|
||||||
|
|
||||||
|
### D7. Token-candidates returns two real fields, two empty objects
|
||||||
|
|
||||||
|
**Choice:** `models` and `modelsWithoutToken` are computed from
|
||||||
|
`model_availability ⋈ account_tokens ⋈ accounts ⋈ sites`.
|
||||||
|
`modelsMissingTokenGroups` and `endpointTypesByModel` are returned as
|
||||||
|
`{}`.
|
||||||
|
|
||||||
|
**Rationale:** Legacy Node behavior depends on
|
||||||
|
`fetchModelPricingCatalog` and `readModelsMarketplaceCache`, neither of
|
||||||
|
which has a Rust counterpart in this tree. Returning empty objects
|
||||||
|
satisfies the schema; the UI degrades gracefully (no missing-group
|
||||||
|
warnings, no per-model endpoint chips). When/if a marketplace cache
|
||||||
|
lands, this handler is the place to wire it in.
|
||||||
|
|
||||||
|
## Risks / Trade-offs
|
||||||
|
|
||||||
|
- **[Stale `model_availability`]** → A rebuild reflects whatever was in
|
||||||
|
the table at run time. Mitigation: the existing
|
||||||
|
`POST /api/accounts/{id}/models/refresh` is the contract for
|
||||||
|
freshness; the auto-rebuild fires after account create, and the
|
||||||
|
manual rebuild button covers re-runs.
|
||||||
|
- **[Long rebuild blocks writes]** → A single transaction holds the
|
||||||
|
write lock for the duration of the rebuild. For tens of accounts and
|
||||||
|
hundreds of models, expected runtime is sub-second. If the install
|
||||||
|
scales past low thousands of channels, switch D2 to per-route
|
||||||
|
transactions or move the rebuild to a background worker.
|
||||||
|
- **[Auto-rebuild noise]** → Fires on every account create, even when
|
||||||
|
`model_availability` for that account is still empty (rebuild becomes
|
||||||
|
a no-op writing nothing). Acceptable: cost is one short transaction
|
||||||
|
with empty desired sets. Logged at `info` so operators can see it.
|
||||||
|
- **[Empty `endpointTypesByModel`]** → Frontend can't show endpoint
|
||||||
|
chips. Cosmetic; not blocking. Documented as a non-goal so
|
||||||
|
follow-ups know where to plug the marketplace cache.
|
||||||
|
- **[`manual_override` semantics depend on `BOOLEAN` storage]** →
|
||||||
|
SQLite stores booleans as 0/1 integers. We decode via `Option<bool>`
|
||||||
|
and `unwrap_or(false)`; rows with `NULL` are treated as auto. Matches
|
||||||
|
schema default.
|
||||||
|
|
||||||
|
## Migration Plan
|
||||||
|
|
||||||
|
No schema changes. Deploy is a code-only swap:
|
||||||
|
|
||||||
|
1. Ship the new binary. The stub handler is replaced atomically; first
|
||||||
|
request after the swap exercises the new pipeline.
|
||||||
|
2. Run `POST /api/routes/rebuild` once after deploy to backfill
|
||||||
|
existing installs whose `model_availability` is already populated.
|
||||||
|
3. New account creations from this point on auto-rebuild.
|
||||||
|
|
||||||
|
**Rollback:** revert the binary. The new code never wrote to a column
|
||||||
|
or table that didn't already exist, and `manual_override` rows are
|
||||||
|
untouched, so rollback restores the stub without data loss.
|
||||||
|
|
||||||
|
## Open Questions
|
||||||
|
|
||||||
|
- Should `route-model-discovery` ship as a separate change, or is
|
||||||
|
"refresh endpoint fills the table" considered the permanent design?
|
||||||
|
Tracked outside this change for now.
|
||||||
|
- Do we want a debounce on the auto-rebuild when batch-creating many
|
||||||
|
accounts in quick succession? Current behavior fires once per
|
||||||
|
`POST /api/accounts` request (covers batch via the multi-token
|
||||||
|
branch), which is fine. If callers ever script many sequential
|
||||||
|
single-account POSTs, consider coalescing.
|
||||||
@@ -0,0 +1,46 @@
|
|||||||
|
## Why
|
||||||
|
|
||||||
|
前端 `/routes` 页面打开后路由列表完全空白，根因是后端 `POST /api/routes/rebuild` 是个空 stub（`src/routes/admin/compat.rs:368-370` 只回 `success: true`），`/api/models/token-candidates` 端点根本没注册。即使站点和账号都配好了，`token_routes` / `route_channels` 这两张表也没人去填，前端只能渲染"暂无路由"。要让路由功能从配好账号到出现在列表里端到端可用，必须把"扫上游模型 → 过滤 → 写表"这条流水线在后端补齐。
|
||||||
|
|
||||||
|
## What Changes
|
||||||
|
|
||||||
|
- **新增**：多平台模型发现（OpenAI 兼容 / Anthropic / 自定义），从站点 `url + api_key + custom_headers` 拉 `/v1/models`，结果按 `(site_id, account_id)` 缓存并尊重 `site_disabled_models` 黑名单。
|
||||||
|
- **新增**：`POST /api/routes/rebuild` 真实实现 — 扫描所有 active 账号，按 `model_name` 聚合为一条 `token_routes` 记录，每个 `(account, source_model)` 落一条 `route_channels` 记录，多账号自动合并到同一 route 下。响应体包含 `rebuild.createdRoutes` / `createdChannels` 以满足 `TokenRoutes.tsx:411-412` 的现有契约。
|
||||||
|
- **新增**：`GET /api/models/token-candidates` 端点 — 返回 `{ models, modelsWithoutToken, modelsMissingTokenGroups, endpointTypesByModel }`，给前端"候选模型"面板和新建路由弹窗的下拉用（`TokenRoutes.tsx:252` 已经在调）。
|
||||||
|
- **新增**：每次"账号成功添加"后由 service 层异步触发一次 rebuild，避免用户必须手点重建按钮。
|
||||||
|
- **修改**：`compat.rs` 中保留的 `rebuild_routes` 占位 handler 删除，迁移到独立 service 模块。
|
||||||
|
- 不破坏已有 `GET /api/routes` / `/lite` / `/summary` 与 `route_channels` 写表 schema。
|
||||||
|
|
||||||
|
## Capabilities
|
||||||
|
|
||||||
|
### New Capabilities
|
||||||
|
|
||||||
|
- `route-model-discovery`: 站点级模型清单的获取、缓存与黑名单过滤。包含多平台 fetcher trait（OpenAI / Anthropic / 自定义 header）、`site_disabled_models` 应用、超时与重试策略、按账号 credential mode（`apikey` / oauth）选取认证头。
|
||||||
|
- `route-rebuild-pipeline`: 把发现到的模型合成路由的写库流水线。负责创建/更新 `token_routes`（按模型名聚合）、为每个 `(account, source_model)` 创建 `route_channels`、保留手工 `manual_override = 1` 通道、回滚已下线模型对应的通道。同步回填 `rebuild.createdRoutes` / `createdChannels` 字段。
|
||||||
|
- `route-candidate-api`: `GET /api/models/token-candidates` 端点。基于 `route-model-discovery` 的缓存返回前端需要的 4 个聚合字段，区分"账号有 token 能用"、"账号在但没 token"、"账号 token 缺组"三种状态。
|
||||||
|
|
||||||
|
### Modified Capabilities
|
||||||
|
|
||||||
|
无（`openspec/specs/` 当前为空，本仓库还没有沉淀过 spec，所以全部走新建）。
|
||||||
|
|
||||||
|
## Impact
|
||||||
|
|
||||||
|
**新增代码：**
|
||||||
|
- `src/services/model_discovery/` — 模型 fetcher trait + 三个 impl（openai / anthropic / generic），缓存层。
|
||||||
|
- `src/services/route_rebuild/` — 重建流水线。
|
||||||
|
- `src/routes/admin/model_candidates.rs` — `GET /api/models/token-candidates` handler。
|
||||||
|
- `src/routes/admin/compat.rs` — 把 `rebuild_routes` 改写为调用新 service。
|
||||||
|
|
||||||
|
**接触的表：**
|
||||||
|
- 读：`sites`, `accounts`, `account_tokens`, `site_disabled_models`
|
||||||
|
- 写：`token_routes`, `route_channels`
|
||||||
|
|
||||||
|
**外部依赖：**
|
||||||
|
- 复用现有 `reqwest` 客户端 + `tokio` runtime，不引入新 crate。
|
||||||
|
- 复用 `state.config.proxy.upstream_timeout_secs` 做模型拉取超时。
|
||||||
|
|
||||||
|
**前端：**无需改动。`TokenRoutes.tsx` 和 `api.ts` 当前调用契约保持不变，列表自然出现。
|
||||||
|
|
||||||
|
**风险：**
|
||||||
|
- 模型扫描走外网，慢账号会拖时间。流水线必须 per-site 并发 + 单 site 超时。
|
||||||
|
- `route_channels` 已有数据时不能粗暴 truncate，需保留 `manual_override = 1` 的人工通道。
|
||||||
+88
@@ -0,0 +1,88 @@
|
|||||||
|
## ADDED Requirements
|
||||||
|
|
||||||
|
### Requirement: `GET /api/models/token-candidates` is registered
|
||||||
|
|
||||||
|
The admin router SHALL register `GET /api/models/token-candidates`
|
||||||
|
behind the standard admin auth layer. The handler SHALL return a JSON
|
||||||
|
object with exactly four top-level keys: `models`,
|
||||||
|
`modelsWithoutToken`, `modelsMissingTokenGroups`, `endpointTypesByModel`.
|
||||||
|
|
||||||
|
#### Scenario: Endpoint returns 200 with the four-key shape
|
||||||
|
- **WHEN** an authenticated admin sends `GET /api/models/token-candidates`
|
||||||
|
- **THEN** the response is HTTP 200 with a JSON body containing keys
|
||||||
|
`models`, `modelsWithoutToken`, `modelsMissingTokenGroups`, and
|
||||||
|
`endpointTypesByModel` — no extra top-level keys, no missing keys
|
||||||
|
|
||||||
|
#### Scenario: Unauthenticated request is rejected
|
||||||
|
- **WHEN** a request without a valid admin token reaches the endpoint
|
||||||
|
- **THEN** the admin auth layer returns 401 before the handler runs
|
||||||
|
|
||||||
|
### Requirement: `models` lists every (model, token) candidate
|
||||||
|
|
||||||
|
The `models` field SHALL be an object whose keys are model names and
|
||||||
|
whose values are arrays of candidate descriptors. A `(model, token)`
|
||||||
|
combination SHALL appear in the `models` array for that model exactly
|
||||||
|
when ALL of the following hold:
|
||||||
|
- a row in `model_availability` reports `available = 1` for the model
|
||||||
|
on the token's account
|
||||||
|
- the token's `enabled` is truthy (NULL treated as enabled)
|
||||||
|
- the token's `value_status` is `'ready'` (NULL treated as ready)
|
||||||
|
- the account's `status = 'active'`
|
||||||
|
- the site's `status` is `'active'` (NULL treated as active)
|
||||||
|
|
||||||
|
Each candidate descriptor SHALL contain: `accountId`, `tokenId`,
|
||||||
|
`tokenName`, `isDefault` (boolean), `username`, `siteId`, `siteName`.
|
||||||
|
Within a single model's array, descriptors SHALL be deduplicated by
|
||||||
|
`tokenId`.
|
||||||
|
|
||||||
|
#### Scenario: Active token with available model appears
|
||||||
|
- **WHEN** account 5 (active, site active) has token 11 (enabled,
|
||||||
|
ready) and `model_availability(account=5, model='gpt-4')` is
|
||||||
|
available
|
||||||
|
- **THEN** `models['gpt-4']` includes a descriptor with
|
||||||
|
`accountId = 5, tokenId = 11`
|
||||||
|
|
||||||
|
#### Scenario: Disabled token is excluded
|
||||||
|
- **WHEN** account 5 has token 11 with `enabled = 0`
|
||||||
|
- **THEN** no descriptor with `tokenId = 11` appears in `models` for
|
||||||
|
any model
|
||||||
|
|
||||||
|
#### Scenario: Duplicate tokens in source data deduped
|
||||||
|
- **WHEN** the underlying join produces two rows for the same
|
||||||
|
`(model, tokenId)` pair (e.g., joins fan out)
|
||||||
|
- **THEN** the descriptor appears exactly once in
|
||||||
|
`models[model_name]`
|
||||||
|
|
||||||
|
### Requirement: `modelsWithoutToken` surfaces accounts missing tokens
|
||||||
|
|
||||||
|
The `modelsWithoutToken` field SHALL be an object keyed by model name.
|
||||||
|
For each `(account, model)` combination where the model is available
|
||||||
|
on the active account but no qualifying token (per the `models`
|
||||||
|
requirement) exists for that account, the handler SHALL include a
|
||||||
|
descriptor `{ accountId, username, siteId, siteName }` under that
|
||||||
|
model name. Descriptors SHALL be deduplicated by `accountId` within
|
||||||
|
each model's array.
|
||||||
|
|
||||||
|
#### Scenario: Account with available model but no enabled tokens
|
||||||
|
- **WHEN** account 8 reports `gpt-4` available but has zero rows in
|
||||||
|
`account_tokens` (or all rows are disabled)
|
||||||
|
- **THEN** `modelsWithoutToken['gpt-4']` includes a descriptor with
|
||||||
|
`accountId = 8`, and that account does NOT appear in `models['gpt-4']`
|
||||||
|
|
||||||
|
#### Scenario: Account with token does not appear here
|
||||||
|
- **WHEN** account 5 reports `gpt-4` available and has a usable token
|
||||||
|
- **THEN** account 5 does NOT appear in `modelsWithoutToken['gpt-4']`
|
||||||
|
|
||||||
|
### Requirement: `modelsMissingTokenGroups` and `endpointTypesByModel` are empty
|
||||||
|
|
||||||
|
The handler SHALL return `modelsMissingTokenGroups` and
|
||||||
|
`endpointTypesByModel` as empty JSON objects (`{}`). The frontend
|
||||||
|
contract permits empty objects and these fields SHALL NOT be
|
||||||
|
populated until a marketplace-cache or pricing-catalog source is
|
||||||
|
added to the backend.
|
||||||
|
|
||||||
|
#### Scenario: Both fields are empty objects
|
||||||
|
- **WHEN** the handler returns successfully under any state of the
|
||||||
|
database
|
||||||
|
- **THEN** `response.modelsMissingTokenGroups === {}` and
|
||||||
|
`response.endpointTypesByModel === {}`
|
||||||
+153
@@ -0,0 +1,153 @@
|
|||||||
|
## ADDED Requirements
|
||||||
|
|
||||||
|
### Requirement: Route rebuild reads from `model_availability`
|
||||||
|
|
||||||
|
The rebuild pipeline SHALL derive its working set from the
|
||||||
|
`model_availability` table joined to `accounts`, filtering out:
|
||||||
|
- accounts whose `status` is not `'active'`
|
||||||
|
- rows whose `available` is not `1` (treating `NULL` as not available)
|
||||||
|
- `(account, model)` combinations covered by `site_disabled_models`
|
||||||
|
for the account's `site_id`
|
||||||
|
|
||||||
|
The pipeline SHALL NOT make outbound HTTP requests to upstream platforms
|
||||||
|
during a rebuild. Refreshing `model_availability` is the responsibility
|
||||||
|
of the existing `POST /api/accounts/{id}/models/refresh` endpoint and
|
||||||
|
is out of scope for this requirement.
|
||||||
|
|
||||||
|
#### Scenario: Inactive account is excluded
|
||||||
|
- **WHEN** an account with `status = 'suspended'` has rows in
|
||||||
|
`model_availability`
|
||||||
|
- **THEN** no `token_routes` row is created on its behalf and no
|
||||||
|
`route_channels` row references its `account_id`
|
||||||
|
|
||||||
|
#### Scenario: Site-disabled model is excluded
|
||||||
|
- **WHEN** `site_disabled_models` contains `(site_id = 7, model_name = 'gpt-x')`
|
||||||
|
and an active account on site 7 has `model_availability` for `gpt-x`
|
||||||
|
- **THEN** the rebuild does not include that account in the channel set
|
||||||
|
for the `gpt-x` route
|
||||||
|
|
||||||
|
### Requirement: Rebuild aggregates one route per distinct model name
|
||||||
|
|
||||||
|
The rebuild SHALL ensure exactly one row exists in `token_routes` for
|
||||||
|
every distinct `model_name` in the filtered working set, keyed by
|
||||||
|
`model_pattern = model_name`. The rebuild SHALL reuse existing rows
|
||||||
|
when their `model_pattern` already matches; new rows SHALL be inserted
|
||||||
|
with `routing_strategy = 'weighted'`, `route_mode = 'pattern'`, and
|
||||||
|
`enabled = 1`.
|
||||||
|
|
||||||
|
The rebuild MUST NOT overwrite `display_name`, `display_icon`, or any
|
||||||
|
other fields on existing `token_routes` rows. Operators rely on these
|
||||||
|
fields being editable from the frontend without being clobbered on
|
||||||
|
re-runs.
|
||||||
|
|
||||||
|
#### Scenario: New model creates a new route
|
||||||
|
- **WHEN** model `claude-4-opus` appears in the working set and no
|
||||||
|
`token_routes` row has `model_pattern = 'claude-4-opus'`
|
||||||
|
- **THEN** the rebuild inserts a `token_routes` row with
|
||||||
|
`model_pattern = 'claude-4-opus'`, `routing_strategy = 'weighted'`,
|
||||||
|
`route_mode = 'pattern'`, `enabled = 1`
|
||||||
|
|
||||||
|
#### Scenario: Existing route is reused without overwrite
|
||||||
|
- **WHEN** a `token_routes` row exists with
|
||||||
|
`model_pattern = 'gpt-4'` and `display_name = 'GPT-4 Turbo (Custom)'`
|
||||||
|
- **THEN** the rebuild leaves `display_name` unchanged and reuses the
|
||||||
|
same `id` for channel reconciliation
|
||||||
|
|
||||||
|
### Requirement: Channel reconciliation preserves manual overrides
|
||||||
|
|
||||||
|
For each route, the rebuild SHALL reconcile `route_channels` rows
|
||||||
|
keyed by `(account_id, source_model)` against the desired set derived
|
||||||
|
from `model_availability`. Rows with `manual_override = 1` SHALL be
|
||||||
|
treated as immutable: not inserted, not updated, not deleted. Rows
|
||||||
|
with `manual_override = 0` (or `NULL`) SHALL be:
|
||||||
|
- inserted when missing from the existing set
|
||||||
|
- hard-deleted when no longer in the desired set
|
||||||
|
|
||||||
|
The desired set SHALL exclude any `(account, model)` combination
|
||||||
|
already covered by a `manual_override = 1` row, so the rebuild never
|
||||||
|
inserts an auto channel that shadows a manual one.
|
||||||
|
|
||||||
|
#### Scenario: Manual channel survives rebuild
|
||||||
|
- **WHEN** a `route_channels` row has `manual_override = 1` and its
|
||||||
|
`(account_id, source_model)` is no longer in the desired set
|
||||||
|
- **THEN** the row is left untouched and counted in
|
||||||
|
`RebuildSummary.skipped_manual_channels`
|
||||||
|
|
||||||
|
#### Scenario: Stale auto channel is deleted
|
||||||
|
- **WHEN** a `route_channels` row has `manual_override = 0`,
|
||||||
|
`account_id = 12`, `source_model = 'old-model'`, and no longer
|
||||||
|
appears in the desired set
|
||||||
|
- **THEN** the row is deleted and counted in
|
||||||
|
`RebuildSummary.deleted_channels`
|
||||||
|
|
||||||
|
#### Scenario: New auto channel is inserted with default token
|
||||||
|
- **WHEN** account 12 newly reports `gpt-5` available and has at least
|
||||||
|
one enabled `account_tokens` row
|
||||||
|
- **THEN** a new `route_channels` row is inserted with `route_id`
|
||||||
|
pointing at the `gpt-5` route, `account_id = 12`, the
|
||||||
|
default-flagged enabled token's id (or the lowest id when none is
|
||||||
|
default), `priority = 0`, `weight = 10`, `enabled = 1`,
|
||||||
|
`manual_override = 0`
|
||||||
|
|
||||||
|
#### Scenario: Account without tokens still gets a channel
|
||||||
|
- **WHEN** account 13 reports `gpt-5` available but has no enabled
|
||||||
|
`account_tokens` rows
|
||||||
|
- **THEN** a `route_channels` row is inserted with `token_id = NULL`,
|
||||||
|
relying on the proxy layer's account-level token fallback
|
||||||
|
|
||||||
|
### Requirement: Rebuild runs in a single transaction
|
||||||
|
|
||||||
|
The rebuild SHALL execute every read and write inside one database
|
||||||
|
transaction, committed at the end. A partial failure MUST NOT leave
|
||||||
|
`token_routes` and `route_channels` mutually inconsistent.
|
||||||
|
|
||||||
|
#### Scenario: Failure mid-rebuild rolls back
|
||||||
|
- **WHEN** an `INSERT` into `route_channels` fails partway through
|
||||||
|
the rebuild
|
||||||
|
- **THEN** the transaction aborts and no `token_routes` insert from
|
||||||
|
the same rebuild persists
|
||||||
|
|
||||||
|
### Requirement: Rebuild summary matches frontend contract
|
||||||
|
|
||||||
|
The rebuild SHALL return a `RebuildSummary` containing the fields
|
||||||
|
`createdRoutes`, `reusedRoutes`, `createdChannels`, `deletedChannels`,
|
||||||
|
`skippedManualChannels`, `scannedModels`, `scannedAccounts` (camelCase
|
||||||
|
when serialized). The `POST /api/routes/rebuild` handler SHALL embed
|
||||||
|
this object under the JSON key `rebuild`, alongside `success: true`.
|
||||||
|
|
||||||
|
`TokenRoutes.tsx:411-412` reads `rebuild.createdRoutes` and
|
||||||
|
`rebuild.createdChannels`; both MUST be present and SHALL count only
|
||||||
|
rows actually inserted by the current rebuild (not pre-existing rows
|
||||||
|
that were reused).
|
||||||
|
|
||||||
|
#### Scenario: Rebuild response shape
|
||||||
|
- **WHEN** the handler completes successfully
|
||||||
|
- **THEN** the response body is
|
||||||
|
`{ "success": true, "rebuild": { "createdRoutes": <n1>,
|
||||||
|
"reusedRoutes": <n2>, "createdChannels": <n3>,
|
||||||
|
"deletedChannels": <n4>, "skippedManualChannels": <n5>,
|
||||||
|
"scannedModels": <n6>, "scannedAccounts": <n7> } }`
|
||||||
|
|
||||||
|
#### Scenario: Reused route does not increment createdRoutes
|
||||||
|
- **WHEN** every model in the working set has a pre-existing
|
||||||
|
`token_routes` row
|
||||||
|
- **THEN** `createdRoutes = 0` and `reusedRoutes` equals the number of
|
||||||
|
distinct models scanned
|
||||||
|
|
||||||
|
### Requirement: Auto-rebuild after account creation
|
||||||
|
|
||||||
|
The system SHALL fire one route rebuild in the background after
|
||||||
|
`POST /api/accounts` returns success for at least one created
|
||||||
|
account. The HTTP response SHALL return without waiting for the
|
||||||
|
rebuild to finish. Errors from the background rebuild SHALL be logged
|
||||||
|
at `warn` level and SHALL NOT affect the response status or body.
|
||||||
|
|
||||||
|
#### Scenario: Single-account create triggers rebuild
|
||||||
|
- **WHEN** `POST /api/accounts` creates exactly one account
|
||||||
|
- **THEN** `spawn_route_rebuild` is invoked once with the cloned
|
||||||
|
`AppState` and the response returns immediately
|
||||||
|
|
||||||
|
#### Scenario: Batch-create with all failures does not trigger
|
||||||
|
- **WHEN** `POST /api/accounts` returns with `created.len() == 0` and
|
||||||
|
`failed.len() > 0`
|
||||||
|
- **THEN** `spawn_route_rebuild` is not invoked
|
||||||
@@ -0,0 +1,52 @@
|
|||||||
|
## 1. Service module
|
||||||
|
|
||||||
|
- [x] 1.1 Create `src/services/route_rebuild.rs` with `RebuildSummary` struct (camelCase serde) and `pub async fn rebuild_routes(&AppState) -> AppResult<RebuildSummary>`
|
||||||
|
- [x] 1.2 Implement transactional pipeline: open `state.db.begin()`, do all reads/writes through `&mut *tx`, commit at the end
|
||||||
|
- [x] 1.3 Query `model_availability ⋈ accounts` filtering by `accounts.status = 'active'`, `available = 1`, and `NOT EXISTS site_disabled_models`
|
||||||
|
- [x] 1.4 Group rows by model name into `HashMap<String, Vec<i64>>` and pre-load existing `(model_pattern, id)` from `token_routes` for reuse
|
||||||
|
- [x] 1.5 Insert new `token_routes` rows with `routing_strategy = 'weighted'`, `route_mode = 'pattern'`, `enabled = 1`; never overwrite existing display fields
|
||||||
|
- [x] 1.6 Per route, partition existing `route_channels` rows into `manual_keys` (skip) vs `auto_by_key` (reconcile)
|
||||||
|
- [x] 1.7 Insert missing auto channels with default-token pick (`ORDER BY COALESCE(is_default,0) DESC, id ASC LIMIT 1`); allow `token_id = NULL` when no tokens exist
|
||||||
|
- [x] 1.8 Hard-delete auto channels that fell out of the desired set; count into `RebuildSummary.deleted_channels`
|
||||||
|
- [x] 1.9 Expose the module from `src/services/mod.rs` (`pub mod route_rebuild;`)
|
||||||
|
|
||||||
|
## 2. Rebuild HTTP handler
|
||||||
|
|
||||||
|
- [x] 2.1 Replace the stub `rebuild_routes()` in `src/routes/admin/compat.rs:368` with a handler that takes `State<AppState>` and calls `services::route_rebuild::rebuild_routes`
|
||||||
|
- [x] 2.2 Return `Json(json!({ "success": true, "rebuild": summary }))` so `TokenRoutes.tsx:411-412` reads `rebuild.createdRoutes` / `createdChannels`
|
||||||
|
|
||||||
|
## 3. Token-candidates handler
|
||||||
|
|
||||||
|
- [x] 3.1 Create `src/routes/admin/model_candidates.rs` with `pub async fn token_candidates(State<AppState>) -> AppResult<Json<Value>>`
|
||||||
|
- [x] 3.2 Query `(model, account, token, site)` join filtering active sites/accounts, enabled tokens, `value_status = 'ready'`; build `models` map deduped by `tokenId` per model
|
||||||
|
- [x] 3.3 Query `(model, account, site)` join for available models on active accounts; diff against the covered set to populate `modelsWithoutToken` deduped by `accountId` per model
|
||||||
|
- [x] 3.4 Return all four keys: `models`, `modelsWithoutToken`, `modelsMissingTokenGroups: {}`, `endpointTypesByModel: {}`
|
||||||
|
- [x] 3.5 Mount `GET /api/models/token-candidates` under the admin auth layer in `src/routes/admin/mod.rs`
|
||||||
|
|
||||||
|
## 4. Auto-rebuild after account create
|
||||||
|
|
||||||
|
- [x] 4.1 Add `fn spawn_route_rebuild(state: AppState)` helper at the bottom of `src/routes/admin/accounts.rs` that `tokio::spawn`s the rebuild and logs the summary at `info`, errors at `warn`
|
||||||
|
- [x] 4.2 Call `spawn_route_rebuild(state.clone())` in the single-account success branch of `create_account`
|
||||||
|
- [x] 4.3 Call `spawn_route_rebuild(state.clone())` in the batch branch when `created.len() > 0`; skip when no accounts were created
|
||||||
|
|
||||||
|
## 5. Verification
|
||||||
|
|
||||||
|
- [x] 5.1 `cargo build` clean — no warnings on the new code
|
||||||
|
- [x] 5.2 Confirm `compat.rs:19` route mount still points at the (now non-stub) `rebuild_routes` handler
|
||||||
|
- [x] 5.3 Confirm `accounts.rs:228` (batch) and `:248` (single) success paths both fire the auto-rebuild
|
||||||
|
|
||||||
|
<!--
|
||||||
|
Descoped (not implemented in this change, see design.md "Non-Goals"):
|
||||||
|
|
||||||
|
- `route-model-discovery` capability (per-platform `/v1/models` fetcher
|
||||||
|
trait with per-(site, account) cache). Deferred per "就用现成的"
|
||||||
|
instruction; existing `POST /api/accounts/{id}/models/refresh` fills
|
||||||
|
`model_availability`, so the pipeline reads from there.
|
||||||
|
- Marketplace pricing-catalog source for `endpointTypesByModel` and
|
||||||
|
`modelsMissingTokenGroups`. Returning `{}` is the documented contract;
|
||||||
|
populate when a Rust-side cache lands.
|
||||||
|
|
||||||
|
These are tracked as potential follow-up changes, not work owed by this
|
||||||
|
change.
|
||||||
|
-->
|
||||||
|
|
||||||
@@ -0,0 +1,2 @@
|
|||||||
|
schema: spec-driven
|
||||||
|
created: 2026-05-19
|
||||||
@@ -0,0 +1,143 @@
|
|||||||
|
## Context
|
||||||
|
|
||||||
|
`src/routes/comfyui_agent.rs` 第 79-91 行的 `upstream_config()` 用 `BOOSAPI_UPSTREAM_BASE` env（默认 `https://api.boos.lat/v1`）+ `BOOSAPI_UPSTREAM_KEY` env（默认 `state.config.auth.proxy_token`）拼出 base/key，三个 handler（chat 第 351 行、images 第 274 行、audio 第 311 行）都直接用 `reqwest::Client` post 到这个外部地址。这导致：
|
||||||
|
|
||||||
|
- ComfyUI Agent 流量绕过本进程已实现的 token_router / capability / proxy_logs / 计费。
|
||||||
|
- 部署必须保证 `api.boos.lat` 可达，否则功能不可用。
|
||||||
|
- 本服务自己就是 OpenAI 兼容网关，"BoosAPI 调远端 BoosAPI" 是反直觉的部署形态。
|
||||||
|
|
||||||
|
数据库现状（`migrations/20260516000001_initial.sql`）：
|
||||||
|
- `downstream_api_keys` 表已存在，`user_id INTEGER REFERENCES users(id)` 可空（无 NOT NULL 约束）。
|
||||||
|
- `key` 列上有 `downstream_api_keys_key_unique` 唯一索引。
|
||||||
|
- `settings` 表是 `(key TEXT PRIMARY KEY, value TEXT)`，可放任意键值。
|
||||||
|
- `find_downstream_key_by_token` 已能按明文 key 查（`src/db/queries/mod.rs:31`）。
|
||||||
|
|
||||||
|
约束：
|
||||||
|
- 保留现栈（axum + sqlx + reqwest），不引入新依赖。
|
||||||
|
- 启动顺序在 `src/main.rs`：load config → `db::create_pool` → migrate → `AppState::new` → `state.comfyui.start()` → `build_router`。Builtin key 引导必须在 migrate 之后、router 之前。
|
||||||
|
- 不破坏既有 `BOOSAPI_UPSTREAM_*` env 在调试场景下的语义。
|
||||||
|
|
||||||
|
## Goals / Non-Goals
|
||||||
|
|
||||||
|
**Goals:**
|
||||||
|
- ComfyUI Agent 默认调用 `http://127.0.0.1:{server.port}/v1/{chat/completions, images/generations, audio/speech}`。
|
||||||
|
- 引入名为 `comfyui-agent` 的内置下游 key，明文持久化到 `settings` 表，重启后保持稳定（不每次换 key）。
|
||||||
|
- 三类调用（chat / images / audio）行为一致：同一 base + 同一 Bearer key + 同一回退策略。
|
||||||
|
- `/api/comfyui-agent/status` 暴露 `mode` 与 `keyName`，供前端排障使用。
|
||||||
|
- 启动时不依赖外部网络。
|
||||||
|
|
||||||
|
**Non-Goals:**
|
||||||
|
- 不为内置 key 增加新的管理界面字段（沿用现有 Downstream Keys 页面，仅做"不可删除"标记）。
|
||||||
|
- 不重写 token_router / capability，本次只是消费它们。
|
||||||
|
- 不做请求级动态 key 切换（每次请求都用同一个 builtin key）。
|
||||||
|
- 不改造 SYSTEM_PROMPT、CHARACTER/SCENE 解析、SSE 转发等业务逻辑。
|
||||||
|
- 不动 `comfyui_process.rs`（嵌入式 ComfyUI 二进制管理）。
|
||||||
|
|
||||||
|
## Decisions
|
||||||
|
|
||||||
|
### 1. 内置 key 用现有 `downstream_api_keys` 表 + `user_id IS NULL` 表示系统级
|
||||||
|
|
||||||
|
**为什么不另建表**：
|
||||||
|
- 想要让 builtin key 与普通用户 key 共享同一套计费、限额、proxy_logs 关联（`usage_logs.downstream_key_id REFERENCES downstream_api_keys(id)` 已存在）。
|
||||||
|
- `user_id` 已经是可空 `INTEGER REFERENCES users(id)`，不需要 schema 改动。
|
||||||
|
|
||||||
|
**Alternative 考虑**：
|
||||||
|
- 单独建一张 `system_api_keys` 表 → 拒绝，要在 `usage_logs` 里加 `system_key_id` 列、改 query，复杂度不成比例。
|
||||||
|
- 复用 `config.auth.proxy_token` 当 Bearer → 拒绝，proxy_token 与 admin auth_token 是不同语义且不在 `downstream_api_keys` 中，无法计费。
|
||||||
|
|
||||||
|
### 2. 密钥明文写到 `settings` 表，key = `builtin_key.comfyui_agent.secret`
|
||||||
|
|
||||||
|
启动引导逻辑（`src/services/builtin_keys.rs::ensure_comfyui_agent_key`）：
|
||||||
|
|
||||||
|
1. `SELECT value FROM settings WHERE key = 'builtin_key.comfyui_agent.secret'`。
|
||||||
|
2. 若有：在 `downstream_api_keys` 中 `find_downstream_key_by_token(value)`，命中则记 id；未命中则进入步骤 3。
|
||||||
|
3. 若无（或步骤 2 未命中）：生成一个 `sk-comfyui-{uuid_v4_no_hyphen}` 形式的 key，先 `INSERT INTO downstream_api_keys(name='comfyui-agent', key=?, user_id=NULL, enabled=1, description='Built-in key for ComfyUI Agent')`，再 `INSERT OR REPLACE INTO settings(key, value)`。
|
||||||
|
4. 返回 `(id, secret)` 给 AppState 缓存（`AppState.builtin_keys.comfyui_agent: (i64, String)`）。
|
||||||
|
|
||||||
|
整个流程在事务中执行，确保 `settings` 与 `downstream_api_keys` 一致。
|
||||||
|
|
||||||
|
**为什么不每次启动重新生成**：
|
||||||
|
- 每次换 key 会让历史 `proxy_logs` / `usage_logs` 仍然引用旧 id，但客户端无法用旧 key 验证。
|
||||||
|
- 重启后 key 稳定也方便用户在前端 Downstream Keys 页面持续看到同一行。
|
||||||
|
|
||||||
|
**为什么不 hash 存储**：
|
||||||
|
- 后端调用方就是后端自己，需要在内存中直接持有明文当 Bearer 用；hash 反而要额外解密通道。
|
||||||
|
- 现有所有 downstream_api_keys 都是明文存储（`find_downstream_key_by_token` 直接 `WHERE key = ?`），保持一致。
|
||||||
|
|
||||||
|
### 3. base 解析优先级：env override → 本地拼接
|
||||||
|
|
||||||
|
```rust
|
||||||
|
fn upstream_config(state: &AppState) -> (String, String, AgentMode) {
|
||||||
|
if let Some(env_base) = std::env::var("BOOSAPI_UPSTREAM_BASE").ok()
|
||||||
|
.map(|v| v.trim().trim_end_matches('/').to_string())
|
||||||
|
.filter(|v| !v.is_empty())
|
||||||
|
{
|
||||||
|
let key = std::env::var("BOOSAPI_UPSTREAM_KEY").ok()
|
||||||
|
.filter(|v| !v.trim().is_empty())
|
||||||
|
.unwrap_or_else(|| state.builtin_keys.comfyui_agent.1.clone());
|
||||||
|
return (env_base, key, AgentMode::FallbackRemote);
|
||||||
|
}
|
||||||
|
|
||||||
|
let host = if state.config.server.host == "0.0.0.0" || state.config.server.host == "::" {
|
||||||
|
"127.0.0.1"
|
||||||
|
} else {
|
||||||
|
&state.config.server.host
|
||||||
|
};
|
||||||
|
let base = format!("http://{}:{}/v1", host, state.config.server.port);
|
||||||
|
let key = state.builtin_keys.comfyui_agent.1.clone();
|
||||||
|
(base, key, AgentMode::Local)
|
||||||
|
}
|
||||||
|
```
|
||||||
|
|
||||||
|
**为什么默认本地而不是默认 env**：用户希望 BoosAPI 自包含运行；env 仅作为"我知道我在干什么"的回退。
|
||||||
|
|
||||||
|
**为什么不 await async client**：base 解析是纯字符串，每次请求重新算成本可忽略（每次请求最多 3 次调用）；不缓存避免与 env 动态变化语义冲突。
|
||||||
|
|
||||||
|
### 4. status 端点扩字段，不改路径
|
||||||
|
|
||||||
|
`GET /api/comfyui-agent/status` 响应增加：
|
||||||
|
|
||||||
|
```json
|
||||||
|
{
|
||||||
|
"success": true,
|
||||||
|
"model": "...",
|
||||||
|
"upstreamBase": "<resolved>",
|
||||||
|
"mode": "local" | "fallback-remote",
|
||||||
|
"keyName": "comfyui-agent",
|
||||||
|
"comfyui": {...}
|
||||||
|
}
|
||||||
|
```
|
||||||
|
|
||||||
|
前端可在标题栏显示 mode 徽标，区分本地 / 回退。
|
||||||
|
|
||||||
|
### 5. 模型名沿用现配置 `state.config.comfyui.agent_model`
|
||||||
|
|
||||||
|
不改 chat 调用的 `model` 字段。如果本地 `/v1/chat/completions` 没有这个模型，会返回 404，让用户去 token_routes 配置；这是正确的失败路径，比 silently fallback 好。
|
||||||
|
|
||||||
|
images 仍用 `gpt-image-2`、audio 仍用 `tts-1` —— 与现状一致，本地 token_router 命中即用。
|
||||||
|
|
||||||
|
## Risks / Trade-offs
|
||||||
|
|
||||||
|
- **[Risk] `127.0.0.1:{port}` 在容器中不可达**：当 server bind `0.0.0.0`，部分容器内 `127.0.0.1` 仍可达 host 进程；但若用户在 Dockerfile 里禁用 loopback，agent 会失败。 → Mitigation：保留 env 回退；status mode 字段帮助排障；文档明确说明。
|
||||||
|
- **[Risk] settings 表在前端被 `app_settings` admin API 暴露**：明文密钥泄漏给 admin。 → Mitigation：在 settings 读写层加 `builtin_key.*` 前缀过滤，列出 settings 时 mask；本次实施任务覆盖该过滤。
|
||||||
|
- **[Risk] 前端 Downstream Keys 页面显示 builtin key 的明文**：现有页面会列出明文。 → Mitigation：在 `routes/user_api_keys.rs`（admin list 入口）里把 `name='comfyui-agent' AND user_id IS NULL` 的 key 标记 `system: true` 并打码 key 显示后 4 位；禁止 DELETE/PATCH。
|
||||||
|
- **[Risk] builtin key 被人为禁用**：admin 在 UI 把 enabled 改为 false，agent 会立即 401。 → Mitigation：启动引导时若发现 enabled=0，主动设回 1 并发 warning；admin API 拒绝对 system key 设 enabled=0。
|
||||||
|
- **[Risk] 引导事务失败导致 settings 记录与 downstream_api_keys 不一致**：例如 settings 插入后进程崩溃，重启时 token 还在 settings 但表里没行。 → Mitigation：步骤 2 已经包含"settings 有值但表没行"的回填路径；幂等。
|
||||||
|
- **[Trade-off] 仍保留 `BOOSAPI_UPSTREAM_*` env**：增加一个分支，但提供排障逃生口；用户态文档明确"仅排障使用"。
|
||||||
|
|
||||||
|
## Migration Plan
|
||||||
|
|
||||||
|
1. **第 1 步**：加 `src/services/builtin_keys.rs` 模块，包含 `ensure_comfyui_agent_key(db: &SqlitePool) -> AppResult<(i64, String)>` 与单测（重复调用幂等、settings 有值但表无行的回填、表内有但 settings 无的恢复）。
|
||||||
|
2. **第 2 步**：在 `AppState` 增加 `builtin_keys: BuiltinKeys { comfyui_agent: (i64, String) }`；`AppState::new` 改为 async（或在 `main.rs` 中先 ensure 再传入）。
|
||||||
|
3. **第 3 步**：重写 `comfyui_agent.rs::upstream_config`，新增 `AgentMode` 枚举；三个 handler 改用新返回值；status 多返回 `mode` / `keyName`。
|
||||||
|
4. **第 4 步**：在 admin 层（`routes/user_api_keys.rs` 或 admin keys 路由）对 `name='comfyui-agent' AND user_id IS NULL` 的行加保护：拒绝 DELETE / 拒绝 enabled=0；list 响应里标 `system=true`。
|
||||||
|
5. **第 5 步**：在 settings 暴露层（确认有没有 generic settings list API）里过滤 / mask `builtin_key.*` 前缀。
|
||||||
|
6. **第 6 步**：手测三类调用 + status mode 切换；`AGENTS.md` / `README.md` 更新。
|
||||||
|
|
||||||
|
回滚：每步独立 commit；如步骤 3 出问题，revert 即可恢复"默认远端"行为。步骤 1-2 是新增代码，无需回滚。
|
||||||
|
|
||||||
|
## Open Questions
|
||||||
|
|
||||||
|
- 是否有 generic `/api/admin/settings` 路由会列出 settings 表全部行？需要 `rg "FROM settings"` 在实施第一步前确认；若有，第 5 步改为对该路由加前缀过滤。
|
||||||
|
- token_router 在没有为 `agent_model` / `gpt-image-2` / `tts-1` 配置路由时的失败模式是 404 还是 500？影响错误文案；实施前在本地起一个空配置环境跑一次。
|
||||||
|
- builtin key 是否需要支持轮换？本次不做，Open Question 记录后续若有需求再开 spec。
|
||||||
@@ -0,0 +1,41 @@
|
|||||||
|
## Why
|
||||||
|
|
||||||
|
`src/routes/comfyui_agent.rs` 目前把 chat / images / audio 三类上游调用全部硬编码走 `https://api.boos.lat/v1`（可被 `BOOSAPI_UPSTREAM_BASE` env 覆盖），完全绕过了本进程已经实现的代理层（token_router、capability、proxy_logs、计费）。这意味着：ComfyUI Agent 的流量看不见、算不上账、用不到本地的智能路由；而本服务自己就是一个 OpenAI 兼容网关，本地 `/v1` 已经能提供 chat/images/audio 三条能力。再继续走外部上游既绕过自家路由又造成"BoosAPI 调 BoosAPI"的奇怪部署形态。
|
||||||
|
|
||||||
|
## What Changes
|
||||||
|
|
||||||
|
- ComfyUI Agent 默认改为调用本进程自身的 `http://127.0.0.1:{server.port}/v1/*`，复用 token_router 与计费体系。
|
||||||
|
- 引入一个**内置下游 API key**（name = `comfyui-agent`，user_id 留空表示系统级），仅 ComfyUI Agent 使用，单独计费/限额；启动时若不存在则自动创建并把明文写入 `app_settings`，重启后复用。
|
||||||
|
- chat / images / audio 三类调用都切到本地 `/v1`，使用上述内置 key 作 Bearer。
|
||||||
|
- 保留 `BOOSAPI_UPSTREAM_BASE` 为**可选回退**：仅当显式设置且非空时，Agent 才直连远端（用于本地路由不可用的排障场景），并在 `/api/comfyui-agent/status` 响应中明确暴露当前生效的 base 与是否为回退。
|
||||||
|
- `/api/comfyui-agent/status` 新增字段 `keyName`（始终是 `comfyui-agent`）、`mode`（`local` 或 `fallback-remote`）。
|
||||||
|
- **BREAKING**（运维侧）：当未设置 `BOOSAPI_UPSTREAM_BASE` 时，部署不再需要外部 `api.boos.lat` 网络可达；但同时要求本进程 `server.port` 已绑定到回环且能被 Agent 自身访问。
|
||||||
|
|
||||||
|
## Capabilities
|
||||||
|
|
||||||
|
### New Capabilities
|
||||||
|
|
||||||
|
- `comfyui-agent-local-routing`: ComfyUI Agent 通过本进程 `/v1` 调用 chat / images / audio，使用专属内置下游 key 计费，并提供 `BOOSAPI_UPSTREAM_BASE` 远端回退。
|
||||||
|
- `builtin-downstream-key`: 系统级（user_id IS NULL）内置下游 API key 的生命周期管理：启动时确保存在、密钥持久化在 `app_settings`、可在 admin API 中可见但不可删除。
|
||||||
|
|
||||||
|
### Modified Capabilities
|
||||||
|
|
||||||
|
<!-- openspec/specs/ 目前为空，没有既有 spec 受影响。 -->
|
||||||
|
|
||||||
|
## Impact
|
||||||
|
|
||||||
|
- **后端代码**
|
||||||
|
- `src/routes/comfyui_agent.rs`：`upstream_config()` 改为从 AppState 拼本地 base + 内置 key；保留 env 回退分支并在响应里暴露 `mode`。
|
||||||
|
- `src/services/`（或新建 `src/services/builtin_keys.rs`）：启动时确保 `comfyui-agent` 内置 key 存在，明文落 `app_settings`。
|
||||||
|
- `src/db/queries/`：新增 `find_or_create_builtin_downstream_key(name)` / `read_builtin_key_secret(name)`。
|
||||||
|
- `src/lib.rs` 的 `AppState::new` 之后、`main.rs` 启动序列中调用 builtin key 引导。
|
||||||
|
- **数据库**
|
||||||
|
- 复用既有 `downstream_api_keys` 表（user_id 设为 NULL 表示系统级）；可能需要一条迁移把 `user_id` 列改为可空（如果已是可空则跳过）。
|
||||||
|
- 复用既有 `app_settings` / `settings` 键值表存放生成的明文密钥（避免每次重启换 key 让历史日志关联断裂）。
|
||||||
|
- **观测**
|
||||||
|
- ComfyUI Agent 流量从此出现在 `proxy_logs` 中，按 key=`comfyui-agent` 可在管理面板筛出。
|
||||||
|
- **配置 / 文档**
|
||||||
|
- `AGENTS.md` 与 `README.md` 的"ComfyUI Agent"段落需要更新：默认本地、env 回退、密钥位置。
|
||||||
|
- **测试**
|
||||||
|
- 新增端到端测试：mock 一个 `/v1/chat/completions` 路径返回固定 SSE，验证 Agent 命中本地。
|
||||||
|
- 新增单测：builtin key 引导 idempotent；status 在 env 设置/未设置时分别返回 `local` / `fallback-remote`。
|
||||||
@@ -0,0 +1,82 @@
|
|||||||
|
## ADDED Requirements
|
||||||
|
|
||||||
|
### Requirement: 启动时确保内置 ComfyUI Agent key 存在
|
||||||
|
|
||||||
|
服务启动序列 SHALL 在数据库迁移完成之后、HTTP 路由对外提供之前，调用 `ensure_comfyui_agent_key`，确保 `downstream_api_keys` 表中存在唯一一行 `name='comfyui-agent' AND user_id IS NULL`，且其明文 key 同步保存在 `settings` 表的 `builtin_key.comfyui_agent.secret` 键。该流程 MUST 幂等。
|
||||||
|
|
||||||
|
#### Scenario: 全新数据库首次启动
|
||||||
|
|
||||||
|
- **WHEN** 启动时数据库无任何 `name='comfyui-agent'` 行，且 settings 表无 `builtin_key.comfyui_agent.secret`
|
||||||
|
- **THEN** 启动结束后两处都存在；密钥形如 `sk-comfyui-<32 位十六进制或等价 uuid>`；`enabled=1`
|
||||||
|
|
||||||
|
#### Scenario: 已有 settings 但缺少表记录（崩溃恢复）
|
||||||
|
|
||||||
|
- **WHEN** settings 表存在 `builtin_key.comfyui_agent.secret = "sk-comfyui-xxx"`，但 `downstream_api_keys` 中没有对应行
|
||||||
|
- **THEN** 引导逻辑用 settings 中的明文回填一行 `downstream_api_keys`，不重新生成密钥
|
||||||
|
|
||||||
|
#### Scenario: 已有表记录但缺少 settings
|
||||||
|
|
||||||
|
- **WHEN** `downstream_api_keys` 中存在 `name='comfyui-agent' AND user_id IS NULL` 行 key=`sk-comfyui-yyy`，但 settings 中无 `builtin_key.comfyui_agent.secret`
|
||||||
|
- **THEN** 引导逻辑把表中的明文写回 settings；密钥保持原值
|
||||||
|
|
||||||
|
#### Scenario: 重复启动不变更密钥
|
||||||
|
|
||||||
|
- **WHEN** 同一数据库连续启动两次
|
||||||
|
- **THEN** 第二次启动后 settings 与表中的 key 与第一次完全一致
|
||||||
|
|
||||||
|
#### Scenario: 引导期间发现 enabled=0 自动恢复
|
||||||
|
|
||||||
|
- **WHEN** 启动时发现已有行 `enabled=0`
|
||||||
|
- **THEN** 引导逻辑把该行 `enabled` 设为 1，并通过 `tracing::warn!` 输出一条提示
|
||||||
|
|
||||||
|
### Requirement: 系统级 key 在管理 API 中受保护
|
||||||
|
|
||||||
|
任何 admin 接口 SHALL 拒绝对 `name='comfyui-agent' AND user_id IS NULL` 的下游 key 执行删除（`DELETE`）或将 `enabled` 设为 false。list / get 类响应 MUST 把这一行标记为 `system: true`，并 mask 明文 key（仅暴露后 4 位）。
|
||||||
|
|
||||||
|
#### Scenario: DELETE 被拒
|
||||||
|
|
||||||
|
- **WHEN** admin 调用 `DELETE` 删除该 key 对应 id
|
||||||
|
- **THEN** 返回 403 / 409 类错误（取实现选择），message 含 `system key` 关键字；数据库中的行不变
|
||||||
|
|
||||||
|
#### Scenario: 禁用被拒
|
||||||
|
|
||||||
|
- **WHEN** admin PATCH/PUT 设 `enabled=false`
|
||||||
|
- **THEN** 返回拒绝错误，行 `enabled` 保持 1
|
||||||
|
|
||||||
|
#### Scenario: 列表响应打码
|
||||||
|
|
||||||
|
- **WHEN** admin GET 列出下游 key
|
||||||
|
- **THEN** 该行字段 `key` 形如 `••••••••<后4位>`，并附 `system: true`
|
||||||
|
|
||||||
|
#### Scenario: 普通 key 不受影响
|
||||||
|
|
||||||
|
- **WHEN** admin 操作其他普通 user_id 非空的 key
|
||||||
|
- **THEN** DELETE / PATCH / list 行为不受系统 key 保护逻辑影响
|
||||||
|
|
||||||
|
### Requirement: settings 中的内置密钥不暴露给 admin 列表
|
||||||
|
|
||||||
|
如果存在 admin 接口列出 `settings` 表内容，该接口 SHALL 过滤所有以 `builtin_key.` 为前缀的 key；按 key 单独读取时，对该前缀返回 404 或带占位符 `"••••"`，不返回实际明文。
|
||||||
|
|
||||||
|
#### Scenario: 列表过滤前缀
|
||||||
|
|
||||||
|
- **WHEN** admin GET `/api/admin/settings`（若存在）
|
||||||
|
- **THEN** 响应中不出现任何 key 以 `builtin_key.` 开头的条目
|
||||||
|
|
||||||
|
#### Scenario: 单读返回占位
|
||||||
|
|
||||||
|
- **WHEN** admin GET `/api/admin/settings/builtin_key.comfyui_agent.secret`（若存在）
|
||||||
|
- **THEN** 返回 404 或 `value: "••••"`，不返回真实明文
|
||||||
|
|
||||||
|
### Requirement: 内置 key 用 user_id IS NULL 与普通 key 区分
|
||||||
|
|
||||||
|
内置下游 key 在 `downstream_api_keys` 表中 SHALL 满足 `user_id IS NULL`；任何普通用户创建的 key MUST 写入非空 `user_id`，确保 `name='comfyui-agent' AND user_id IS NULL` 的查询能稳定匹配仅一行。
|
||||||
|
|
||||||
|
#### Scenario: 普通用户不能创建同名 key 占用系统标识
|
||||||
|
|
||||||
|
- **WHEN** 普通用户调用创建 key 接口，name 设为 `comfyui-agent`
|
||||||
|
- **THEN** 该 key 仍写入非空 user_id；系统级查询不会返回该行
|
||||||
|
|
||||||
|
#### Scenario: 引导查询仅匹配系统行
|
||||||
|
|
||||||
|
- **WHEN** 数据库中同时存在 user_id=5 且 name='comfyui-agent' 的普通 key 与 user_id IS NULL 的系统 key
|
||||||
|
- **THEN** `ensure_comfyui_agent_key` 仅作用于 user_id IS NULL 的那一行
|
||||||
+81
@@ -0,0 +1,81 @@
|
|||||||
|
## ADDED Requirements
|
||||||
|
|
||||||
|
### Requirement: ComfyUI Agent 默认调用本进程 /v1
|
||||||
|
|
||||||
|
ComfyUI Agent 的三类上游调用（`/v1/chat/completions`、`/v1/images/generations`、`/v1/audio/speech`）SHALL 默认指向本进程自身，base 由 `state.config.server.host` 与 `state.config.server.port` 拼出（host 为 `0.0.0.0` 或 `::` 时改写为 `127.0.0.1`）。Agent MUST NOT 在默认配置下访问任何外部主机。
|
||||||
|
|
||||||
|
#### Scenario: 默认 base 指向本机回环
|
||||||
|
|
||||||
|
- **WHEN** `BOOSAPI_UPSTREAM_BASE` 未设置且 `server.port = 4000`、`server.host = 0.0.0.0`
|
||||||
|
- **THEN** chat / images / audio 三类调用的 base 解析为 `http://127.0.0.1:4000/v1`
|
||||||
|
|
||||||
|
#### Scenario: 显式 host 被保留
|
||||||
|
|
||||||
|
- **WHEN** `server.host = 10.0.0.5`、`server.port = 4000` 且未设置 env
|
||||||
|
- **THEN** base 解析为 `http://10.0.0.5:4000/v1`，不被改写为回环
|
||||||
|
|
||||||
|
#### Scenario: 三类调用使用同一 base
|
||||||
|
|
||||||
|
- **WHEN** Agent 在同一次会话中触发 chat、images、audio 三类调用
|
||||||
|
- **THEN** 三次请求的 base 完全一致（来自同一 `upstream_config` 调用结果）
|
||||||
|
|
||||||
|
### Requirement: env 显式设置时回退远端
|
||||||
|
|
||||||
|
当且仅当 `BOOSAPI_UPSTREAM_BASE` 设置为非空字符串时，Agent SHALL 改用该 env 值作为 base，并使用 `BOOSAPI_UPSTREAM_KEY`（若设置且非空）或内置 key 的明文作为 Bearer。该分支以下统称 `fallback-remote` 模式。
|
||||||
|
|
||||||
|
#### Scenario: env 非空触发回退
|
||||||
|
|
||||||
|
- **WHEN** `BOOSAPI_UPSTREAM_BASE = "https://api.boos.lat/v1"`
|
||||||
|
- **THEN** 三类调用的 base 解析为 `https://api.boos.lat/v1`，且 `mode = "fallback-remote"`
|
||||||
|
|
||||||
|
#### Scenario: env 为空字符串视同未设置
|
||||||
|
|
||||||
|
- **WHEN** `BOOSAPI_UPSTREAM_BASE = ""` 或 `BOOSAPI_UPSTREAM_BASE = " "`
|
||||||
|
- **THEN** base 解析回本机回环，`mode = "local"`
|
||||||
|
|
||||||
|
#### Scenario: env key 缺省时复用内置 key
|
||||||
|
|
||||||
|
- **WHEN** 设置了 `BOOSAPI_UPSTREAM_BASE` 但未设置 `BOOSAPI_UPSTREAM_KEY`
|
||||||
|
- **THEN** Bearer 使用 `comfyui-agent` 内置 key 的明文，不退化为 `proxy_token`
|
||||||
|
|
||||||
|
### Requirement: 本地调用使用内置下游 key 计费
|
||||||
|
|
||||||
|
`mode = "local"` 时，Agent SHALL 在请求头 `Authorization` 中携带 `Bearer <comfyui-agent 内置 key 明文>`。该 key MUST 与本进程 `downstream_api_keys` 表中 `name='comfyui-agent' AND user_id IS NULL` 的行一致，使所有 Agent 流量在 `usage_logs` 中以该 key 归集。
|
||||||
|
|
||||||
|
#### Scenario: chat 请求携带内置 key
|
||||||
|
|
||||||
|
- **WHEN** Agent 在本地模式下发起 `/v1/chat/completions`
|
||||||
|
- **THEN** 请求头 `Authorization` 形如 `Bearer sk-comfyui-...`，且该值等于 settings 表中 `builtin_key.comfyui_agent.secret` 的值
|
||||||
|
|
||||||
|
#### Scenario: usage_logs 归集到内置 key
|
||||||
|
|
||||||
|
- **WHEN** 一次本地模式下的 chat 请求成功
|
||||||
|
- **THEN** `usage_logs` 中存在一条记录，`downstream_key_id` 等于内置 key 在 `downstream_api_keys` 中的 id
|
||||||
|
|
||||||
|
### Requirement: status 端点暴露生效模式
|
||||||
|
|
||||||
|
`GET /api/comfyui-agent/status` 响应 SHALL 在原有字段基础上增加 `mode: "local" | "fallback-remote"` 与 `keyName: "comfyui-agent"` 两个字段；`upstreamBase` 字段 MUST 反映当前实际生效的 base（已 trim 末尾斜杠）。
|
||||||
|
|
||||||
|
#### Scenario: 默认配置返回 local
|
||||||
|
|
||||||
|
- **WHEN** 客户端在未设置 env 的情况下 GET `/api/comfyui-agent/status`
|
||||||
|
- **THEN** 响应包含 `"mode": "local"`、`"keyName": "comfyui-agent"`、`upstreamBase` 形如 `http://127.0.0.1:<port>/v1`
|
||||||
|
|
||||||
|
#### Scenario: env 设置时返回 fallback-remote
|
||||||
|
|
||||||
|
- **WHEN** `BOOSAPI_UPSTREAM_BASE = "https://api.boos.lat/v1"` 时 GET 该端点
|
||||||
|
- **THEN** 响应 `"mode": "fallback-remote"`、`upstreamBase = "https://api.boos.lat/v1"`
|
||||||
|
|
||||||
|
#### Scenario: 鉴权失败仍不暴露密钥
|
||||||
|
|
||||||
|
- **WHEN** 客户端用错误的 admin token 请求
|
||||||
|
- **THEN** 返回 401，且响应体不包含 `keyName` / `upstreamBase` 任一字段
|
||||||
|
|
||||||
|
### Requirement: 本地路由失败必须如实返回
|
||||||
|
|
||||||
|
当本地 `/v1` 没有配置对应模型路由（chat 的 `agent_model`、images 的 `gpt-image-2`、audio 的 `tts-1`），Agent SHALL 把上游返回的错误状态码 / 文案如实通过 SSE `error` 事件转发给前端，MUST NOT 在收到 4xx 时自动切到 `fallback-remote`。
|
||||||
|
|
||||||
|
#### Scenario: 本地缺少 agent_model 路由
|
||||||
|
|
||||||
|
- **WHEN** 本地模式下发起 chat，token_router 返回 404
|
||||||
|
- **THEN** SSE 流中出现 `error` 事件，message 包含上游 404 内容；后续请求仍走本地，不自动切远端
|
||||||
@@ -0,0 +1,85 @@
|
|||||||
|
## 1. 内置 key 引导逻辑
|
||||||
|
|
||||||
|
- [x] 1.1 新建 `src/services/builtin_keys.rs`，定义 `pub struct BuiltinKeys { pub comfyui_agent: (i64, String) }` 与 `pub async fn ensure_comfyui_agent_key(db: &SqlitePool) -> AppResult<(i64, String)>`
|
||||||
|
- [x] 1.2 实现 `ensure_comfyui_agent_key`：在事务中执行查 settings → 查 downstream_api_keys → 缺什么补什么，确保两边一致；密钥格式 `sk-comfyui-{uuid_v4 去横线}`，settings key = `builtin_key.comfyui_agent.secret`
|
||||||
|
- [x] 1.3 引导逻辑发现 `enabled=0` 时主动设回 1，并 `tracing::warn!("comfyui-agent builtin key was disabled; re-enabling")`
|
||||||
|
- [x] 1.4 在 `src/services/mod.rs` 中 `pub mod builtin_keys;`
|
||||||
|
- [x] 1.5 单测：用 `tempfile::NamedTempFile` + 内存 sqlite 覆盖 5 个场景（首次启动、settings 有表无、表有 settings 无、重复启动密钥不变、enabled=0 自动恢复）
|
||||||
|
|
||||||
|
## 2. AppState 接入
|
||||||
|
|
||||||
|
- [x] 2.1 在 `src/lib.rs` 中给 `AppState` 增加 `pub builtin_keys: Arc<services::builtin_keys::BuiltinKeys>`
|
||||||
|
- [x] 2.2 修改 `AppState::new` 签名为 async，或在 `main.rs` 中改为先 `ensure_comfyui_agent_key` 拿到结果再调 `AppState::new(config, pool, builtin_keys)`
|
||||||
|
- [x] 2.3 调整 `src/main.rs` 启动序列：migrate 之后、`AppState::new` 之前调用引导；任何 builtin key 相关错误用 `anyhow::Context` 包裹后向上传播，不让服务在没有 builtin key 时启动
|
||||||
|
- [x] 2.4 跑 `cargo build`、`cargo test --lib services::builtin_keys`，确认无回归
|
||||||
|
|
||||||
|
## 3. comfyui_agent.rs 切换调用目标
|
||||||
|
|
||||||
|
- [x] 3.1 在 `src/routes/comfyui_agent.rs` 顶部加 `enum AgentMode { Local, FallbackRemote }`，派生 Serialize 以便嵌入 JSON
|
||||||
|
- [x] 3.2 重写 `upstream_config(state) -> (String, String, AgentMode)`：env 非空（trim 后）走 fallback-remote 分支并把 `BOOSAPI_UPSTREAM_KEY` 缺省时用 builtin key；否则用 `state.config.server.host/port` 拼 `http://{host}/v1`，host 是 `0.0.0.0` 或 `::` 时改写 `127.0.0.1`
|
||||||
|
- [x] 3.3 三个 handler（chat 第 351 行、images 第 274 行、audio 第 311 行）全部改为接受 `(api_base, api_key, _mode)`，bearer 用新 `api_key`
|
||||||
|
- [x] 3.4 `status_handler` 响应增加 `"mode": <local|fallback-remote>`、`"keyName": "comfyui-agent"`，`upstreamBase` 仍为已 trim 的字符串
|
||||||
|
- [x] 3.5 单测（用 `axum-test`）：env 未设、env 设、env 空字符串三种条件下 `/api/comfyui-agent/status` 返回的 mode 与 upstreamBase 正确
|
||||||
|
- [x] 3.6 集成测：`wiremock` mock 一个 `/v1/chat/completions` 返回固定 SSE，断言 Agent 在 local 模式下命中该 mock
|
||||||
|
|
||||||
|
## 4. admin 层保护内置 key
|
||||||
|
|
||||||
|
- [x] 4.1 用 `rg "FROM downstream_api_keys" /root/metapi/src/routes` 找到 admin list / get / patch / delete 入口
|
||||||
|
- [x] 4.2 在 list 响应映射处过滤：若 `name='comfyui-agent' AND user_id IS NULL` 则附加 `system: true` 字段、把 `key` 改为 `••••••••<后4位>`
|
||||||
|
- [x] 4.3 在 DELETE 入口加守卫：命中系统 key 返回 403 + message `system key cannot be deleted`
|
||||||
|
- [x] 4.4 在 PATCH/PUT 入口加守卫：当请求体试图把系统 key 的 `enabled` 设为 false 时拒绝
|
||||||
|
- [x] 4.5 路由级单测：3 条用例覆盖打码 / DELETE 拒绝 / 禁用拒绝；普通 key 操作不受影响 1 条用例
|
||||||
|
|
||||||
|
## 5. settings 暴露面过滤
|
||||||
|
|
||||||
|
- [x] 5.1 用 `rg "FROM settings"` 与 `rg "/api/admin/settings"` 找到 settings list / get 入口（如有）
|
||||||
|
- [x] 5.2 list 入口过滤所有以 `builtin_key.` 为前缀的 key
|
||||||
|
- [x] 5.3 单 key 读入口：以 `builtin_key.` 开头时返回 404 或 `value: "••••"`，不返回真实明文
|
||||||
|
- [x] 5.4 若上述路由不存在，把当前调研结果记入 commit message，跳过本组任务但保留 spec scenarios 待将来实现兜底
|
||||||
|
|
||||||
|
> 调研结论：当前 `src/routes/admin/mod.rs` 仅暴露 `GET /api/settings` (list-all) 与 `PUT /api/settings` (upsert)，没有单 key GET。
|
||||||
|
> 已实现：
|
||||||
|
> - `get_settings` 在 `BTreeMap` 收敛前过滤 `key.starts_with("builtin_key.")`
|
||||||
|
> - `update_setting` 对 `builtin_key.*` 写入返回 403 `system setting cannot be modified`，作为防御性兜底（即便 UI 永远看不到该 key，也防止已知 key 名被直接 PUT 篡改）
|
||||||
|
> - 单 key GET 路由不存在，无需兜底；若将来新增需保留同样的 prefix 守卫
|
||||||
|
|
||||||
|
## 6. 文档与冒烟
|
||||||
|
|
||||||
|
- [x] 6.1 更新 `AGENTS.md` 第 32-36 行的 ComfyUI Agent 段落：默认本地、env 回退语义、密钥位置（settings + downstream_api_keys）
|
||||||
|
- [x] 6.2 更新 `README.md` 中相关段落（搜索"ComfyUI Agent"出现处）
|
||||||
|
- [ ] 6.3 跑全量校验：`cargo test`、`cargo clippy --all-targets`
|
||||||
|
- [ ] 6.4 手测：默认启动一次 chat / images / audio 各一次；设置 `BOOSAPI_UPSTREAM_BASE=https://api.boos.lat/v1` 再启动一次三类调用；两次 status 响应 mode 字段切换正确（用户手测，等待执行）
|
||||||
|
- [x] 6.5 在 PR 描述中列出每条 spec scenario 的对应测试位置，便于 reviewer 校对（mapping 见下文「附：scenario → test 映射」）
|
||||||
|
|
||||||
|
## 附：scenario → test 映射
|
||||||
|
|
||||||
|
### specs/builtin-downstream-key/spec.md
|
||||||
|
|
||||||
|
- 启动时确保内置 key 存在
|
||||||
|
- 全新数据库首次启动 → `src/services/builtin_keys.rs` tests::ensures_key_on_first_boot
|
||||||
|
- 已有 settings 缺表记录（崩溃恢复） → `src/services/builtin_keys.rs` tests::backfills_table_from_settings
|
||||||
|
- 已有表记录缺 settings → `src/services/builtin_keys.rs` tests::backfills_settings_from_table
|
||||||
|
- 重复启动不变更密钥 → `src/services/builtin_keys.rs` tests::idempotent_across_two_boots
|
||||||
|
- enabled=0 自动恢复 → `src/services/builtin_keys.rs` tests::recovers_disabled_row
|
||||||
|
- 系统级 key 在管理 API 中受保护
|
||||||
|
- DELETE 被拒 / 禁用被拒 / 列表打码 / 普通 key 不受影响 → `src/routes/admin/downstream_keys.rs` `#[cfg(test)] mod tests`
|
||||||
|
- settings 中的内置密钥不暴露给 admin 列表
|
||||||
|
- 列表过滤前缀 → `src/routes/admin/mod.rs` tests::settings_list_filters_builtin_keys
|
||||||
|
- 单读返回占位 → 单 key GET 路由不存在；防御性写入守卫由 tests::settings_update_rejects_builtin_key 覆盖（详见 5.4 调研结论）
|
||||||
|
- 内置 key 用 user_id IS NULL 与普通 key 区分
|
||||||
|
- 普通用户不能创建同名 key 占用系统标识 / 引导查询仅匹配系统行 → builtin_keys.rs 引导逻辑使用 `name='comfyui-agent' AND user_id IS NULL` 限定 + admin tests 覆盖
|
||||||
|
|
||||||
|
### specs/comfyui-agent-local-routing/spec.md
|
||||||
|
|
||||||
|
- ComfyUI Agent 默认调用本进程 /v1
|
||||||
|
- 默认 base 指向本机回环 / 显式 host 被保留 / 三类调用使用同一 base → `src/routes/comfyui_agent.rs` tests::status_reflects_env_state_across_three_scenarios（Scenario A 覆盖回环改写、Scenario B 覆盖显式 host、共享 upstream_config 即覆盖三类一致）
|
||||||
|
- env 显式设置时回退远端
|
||||||
|
- env 非空触发回退 / env 空字符串视同未设置 → tests::status_reflects_env_state_across_three_scenarios（Scenario B + C）
|
||||||
|
- env key 缺省时复用内置 key → upstream_config 实现内联（见 routes/comfyui_agent.rs 的 fallback-remote 分支）
|
||||||
|
- 本地调用使用内置下游 key 计费
|
||||||
|
- chat 请求携带内置 key → tests::agent_local_mode_routes_chat_to_local_upstream（wiremock 断言 `Bearer sk-comfyui-test`）
|
||||||
|
- usage_logs 归集 → 间接覆盖（agent 经本进程 /v1，途经 token_router 自然写入 usage_logs，由现有 routes/proxy 路径覆盖）
|
||||||
|
- status 端点暴露生效模式
|
||||||
|
- 默认配置返回 local / env 设置时返回 fallback-remote → tests::status_reflects_env_state_across_three_scenarios
|
||||||
|
- 鉴权失败仍不暴露密钥 → 由 status_handler 鉴权中间件覆盖（无新增分支）
|
||||||
|
- 本地路由失败必须如实返回 → 由 SSE error 事件实现保留原行为，未新增专项测；如后续要硬覆盖，建议加 wiremock 404 + 断言 SSE error 事件
|
||||||
@@ -0,0 +1,2 @@
|
|||||||
|
schema: spec-driven
|
||||||
|
created: 2026-05-20
|
||||||
@@ -0,0 +1,102 @@
|
|||||||
|
## Context
|
||||||
|
|
||||||
|
Today the Rust backend (`bosapi v1.3.0`) spawns a single embedded ComfyUI subprocess at boot via `services/comfyui_process.rs::ComfyUIProcess::new(...)` and stores it in `AppState.comfyui: Arc<ComfyUIProcess>`. All `/api/video-agent/*` handlers and any future `/comfy/*` traffic share this one process — so workflows, queue items, custom node state, and uploaded media leak between users. We need to support 100+ users with full per-user isolation, but we have no GPU (we use the BoosAPI text/image/TTS/video endpoints exclusively), so per-user instances are cheap (~1–2 GB resident with the API-only profile).
|
||||||
|
|
||||||
|
The user already runs ComfyUI in Docker as image `bosapi-comfyui:ready` and has accepted the "立刻可用 + 完全隔离" path: per-user on-demand container, with K3S as the v2 target. They added two late constraints we must honor:
|
||||||
|
|
||||||
|
1. Cap each container at **2 CPU / 2 GB RAM** (`--cpus=2 --memory=2g --memory-swap=2g`).
|
||||||
|
2. The existing `/api/video-agent/*` route surface is the in-app ComfyUI consumer; it must keep working for the same frontend with no contract change.
|
||||||
|
|
||||||
|
Auth is the existing JWT cookie. Users are identified by `user_id` from the cookie's claims. There is no per-org limit yet — only a single global cap.
|
||||||
|
|
||||||
|
## Goals / Non-Goals
|
||||||
|
|
||||||
|
**Goals:**
|
||||||
|
- Per-user ComfyUI isolation: workflows, queue, history, and uploaded files MUST NOT cross users.
|
||||||
|
- Lazy startup: a container exists only when the user is actively working; reaped after **60 minutes** idle.
|
||||||
|
- Hard ceiling: at most **30 concurrent containers**, hard-capped to bound memory at ~60 GB.
|
||||||
|
- Per-container resource cap: **2 CPU / 2 GB RAM** (memory + swap).
|
||||||
|
- Persistent custom workflows and outputs at `/data/comfyui-users/<user_id>/`.
|
||||||
|
- Reuse the existing `/api/video-agent/*` route surface unchanged from the frontend's perspective.
|
||||||
|
- New `/comfy/*` HTTP+WebSocket reverse proxy with cookie auth (no username in URL).
|
||||||
|
- Be K3S-portable: the container backend lives behind a `ContainerRuntime` trait so v2 can swap Docker for K3S with no caller changes.
|
||||||
|
- Built-in ComfyUI OpenAI nodes work via the `BOOSAPI_BUILTIN_KEY` env-var fallback (already patched in `_helpers.py`; bake into image).
|
||||||
|
|
||||||
|
**Non-Goals:**
|
||||||
|
- GPU support, local model weights, training. We are API-only.
|
||||||
|
- Per-organization quotas or per-user concurrency budgets — single global cap of 30 only.
|
||||||
|
- K3S backend implementation. v1 ships Docker only; the trait is the only K3S investment.
|
||||||
|
- Migrating away from `bosapi_comfy_nodes.py`. It coexists with the patched internal nodes.
|
||||||
|
- Replacing the existing dev-time `docker-compose.comfyui.yml`. Production stops using it; dev keeps it.
|
||||||
|
|
||||||
|
## Decisions
|
||||||
|
|
||||||
|
### 1. Per-user on-demand Docker container (vs shared+key-injection or per-node-key)
|
||||||
|
**Chosen:** spawn `bosapi-comfyui:multi-tenant` once per active user via `docker run`. **Rationale:** complete filesystem and queue isolation; no need to rewrite ComfyUI's WS protocol to mux users; the user's own framing was "立刻可用 + 完全隔离 ... 后期用K3S". **Alternatives considered:** (a) shared instance with per-request key injection — leaks queue/history across users, requires a reverse-engineered WS proxy that rewrites `client_id`, fragile; (b) per-node API-key — loses access to ComfyUI's built-in OpenAI nodes and prebuilt templates the user explicitly wants kept ("内置的节点和模板都很好用").
|
||||||
|
|
||||||
|
### 2. Pool sizing: 30 hard cap, 60-minute idle reap
|
||||||
|
**Chosen:** `MAX_CONCURRENT=30`, `IDLE_TIMEOUT=60min`. **Rationale:** 30 × 2 GB = 60 GB resident, fits a single mid-size box; 60 min covers a typical creative session without forcing repeated cold starts. Beyond 30, the pool returns `503` with `Retry-After`. **Alternative considered:** unbounded with LRU evict — but evicting a busy queue mid-render is worse UX than asking the next user to wait.
|
||||||
|
|
||||||
|
### 3. URL routing: `/comfy/*` at host root, identity from JWT cookie
|
||||||
|
**Chosen:** mount a reverse-proxy router at `/comfy/*` that reads the JWT cookie, resolves the user's container, and proxies HTTP+WS to `127.0.0.1:<port>`. **Rationale:** ComfyUI's HTML/JS asset paths assume root-relative URLs and break under subpath rewriting; putting the user_id in the URL path leaks it and complicates sharing/auth. **Alternative considered:** subdomain-per-user (`<user>.comfy.boos.lat`) — needs wildcard TLS and DNS, overkill for v1.
|
||||||
|
|
||||||
|
### 4. Container backend behind `ContainerRuntime` trait
|
||||||
|
**Chosen:**
|
||||||
|
```rust
|
||||||
|
#[async_trait]
|
||||||
|
pub trait ContainerRuntime: Send + Sync {
|
||||||
|
async fn start(&self, spec: ContainerSpec) -> Result<ContainerHandle>;
|
||||||
|
async fn stop(&self, id: &str) -> Result<()>;
|
||||||
|
async fn is_alive(&self, id: &str) -> Result<bool>;
|
||||||
|
}
|
||||||
|
```
|
||||||
|
v1 ships `DockerRuntime` (shells out to `docker` or talks to `/var/run/docker.sock`). v2 adds `K3sRuntime`. **Rationale:** the user has already picked K3S as v2; constraining now costs little and saves a rewrite. **Alternative considered:** hard-code `bollard` calls into the pool — couples lifecycle to Docker, rejected.
|
||||||
|
|
||||||
|
### 5. Resource limits per container
|
||||||
|
**Chosen:** `--cpus=2 --memory=2g --memory-swap=2g`, plus `--pids-limit=512`, `--read-only` rootfs with `--tmpfs /tmp:size=512m`, `--cap-drop=ALL`, `--security-opt=no-new-privileges`. Volume mounts: `/data/comfyui-users/<user_id>/input`, `/output`, `/temp`, `/user`, `/custom_nodes_user`. **Rationale:** user explicitly said "Docker的CPU RAM都要做限制 给个2C2G". The hardening flags follow defense-in-depth; ComfyUI does not need privileges or rootfs writes beyond the mounted volumes.
|
||||||
|
|
||||||
|
### 6. Idle reaper as a background tokio task
|
||||||
|
**Chosen:** a single tokio task ticks every 60 s, scans `pool.containers`, and stops any whose `last_used_at` is > 60 min ago. State recorded in-memory; a crash drops it but containers self-stop on `--rm` exit anyway (we do `docker stop`, not `docker kill`, so ComfyUI flushes). **Alternative considered:** Docker's own `--stop-timeout` + cron — split control plane between Rust and shell, harder to reason about.
|
||||||
|
|
||||||
|
### 7. Refactor of `/api/video-agent/*` (preserve frontend contract)
|
||||||
|
**Chosen:** every handler in `routes/video_agent.rs` extracts `user_id` from the JWT cookie, calls `state.comfyui_pool.get_or_start(user_id).await?`, and forwards to that handle's `start/stop/restart/status` or HTTP client. Status responses gain `containerState: "starting" | "ready" | "stopped"` so the frontend can render the new full-screen loading overlay. The existing JSON keys remain. **Alternative considered:** introduce a v2 namespace `/api/v2/video-agent/*` — gratuitous churn on the frontend.
|
||||||
|
|
||||||
|
### 8. AppState rename
|
||||||
|
**Chosen:** `pub comfyui_pool: Arc<ComfyUIPool>` replaces `pub comfyui: Arc<ComfyUIProcess>`. The old `ComfyUIProcess` type is **kept** — it's now the implementation detail used inside `DockerRuntime` to model the in-container ComfyUI subprocess (start/stop, port readiness probe, log buffer). `lib.rs` no longer constructs it directly; `DockerRuntime::start` does.
|
||||||
|
|
||||||
|
### 9. Image build
|
||||||
|
**Chosen:** new `docker/Dockerfile.comfyui` based on the user's existing `bosapi-comfyui:ready` lineage. Adds:
|
||||||
|
- the patched `comfy_source/comfy_api_nodes/util/_helpers.py` (env-var fallback)
|
||||||
|
- `bosapi_comfy_nodes.py` baked in (no host bind-mount, so it survives `--read-only`)
|
||||||
|
- `ENTRYPOINT` reads `BOOSAPI_BUILTIN_KEY`, `BOOSAPI_BASE`, `BOOSAPI_KEY`, then exec's `python main.py --listen 0.0.0.0 --port 8188 --comfy-api-base "$BOOSAPI_BASE/api/comfyui-agent"`.
|
||||||
|
Tag: `bosapi-comfyui:multi-tenant`. `build-and-run.sh` builds this image; the runtime spawns containers from it on demand.
|
||||||
|
|
||||||
|
## Risks / Trade-offs
|
||||||
|
|
||||||
|
- **Cold-start latency 10–20 s** → mitigated by full-screen loading UI with progress text; pool exposes `containerState` so the frontend can poll.
|
||||||
|
- **30-container ceiling can reject users at peak** → 503 with `Retry-After: 60`; UI shows a friendly "服务繁忙" page. Monitor and raise the cap (or move to K3S) if it bites.
|
||||||
|
- **Docker socket exposure** if the API runs in Docker → mount `/var/run/docker.sock` only, restrict to root in API container, never expose to user containers. Document for ops.
|
||||||
|
- **Port pool exhaustion (8200–8299 = 100 slots, cap is 30)** → 70-port headroom is fine; allocator picks lowest-free; on container stop the port is released.
|
||||||
|
- **Idle reap mid-render** → reap only when no active WS connections AND `last_used_at` > 60 min; bump `last_used_at` on every HTTP+WS frame. Acceptable false-negative rate.
|
||||||
|
- **Persistent volume disk growth** → no quota in v1. Document; add a per-user disk cap (e.g. 5 GB via `du`+reject) only if it becomes a real problem.
|
||||||
|
- **Built-in OpenAI nodes auth** depends on `BOOSAPI_BUILTIN_KEY` env reaching the container → enforce via `ContainerSpec.env` in the pool, not opt-in.
|
||||||
|
- **K3S abstraction tax** → only the trait shape; we still have to design pod specs later. Acceptable.
|
||||||
|
- **State loss on API restart** → containers keep running but the pool's in-memory map is gone. Mitigation: on startup, list `comfyui-user-*` containers, reconcile map. Tracked as a task.
|
||||||
|
|
||||||
|
## Migration Plan
|
||||||
|
|
||||||
|
1. Land `ContainerRuntime` trait + `DockerRuntime` + `ComfyUIPool` behind a feature flag `multi_tenant_comfy` (default off). The old `state.comfyui` singleton stays intact.
|
||||||
|
2. Build `bosapi-comfyui:multi-tenant` image; smoke-test by hand with `docker run`.
|
||||||
|
3. Add `/comfy/*` proxy and a single new `state.comfyui_pool` field alongside `state.comfyui`.
|
||||||
|
4. Refactor `/api/video-agent/*` handlers to read from `state.comfyui_pool` when the flag is on, else fall back to `state.comfyui`.
|
||||||
|
5. Toggle the flag on in staging; verify with two real users that workflows/uploads do not cross.
|
||||||
|
6. Remove `state.comfyui` and the flag. `docker-compose.comfyui.yml` becomes dev-only.
|
||||||
|
7. Document host requirements in `README` (Docker socket, `/data/comfyui-users/`, image build).
|
||||||
|
|
||||||
|
**Rollback:** flip the feature flag off and redeploy. Containers from the multi-tenant path keep running until idle-reaped; no data loss because volumes are persistent.
|
||||||
|
|
||||||
|
## Open Questions
|
||||||
|
|
||||||
|
- Should the API container run with `docker:dind` or mount the host socket? Default to host socket; revisit if ops needs isolation.
|
||||||
|
- Where do we publish the loading-state UI string set? Put it next to other i18n strings in the frontend repo.
|
||||||
|
- Do we expose `/comfy/admin/*` for ops to introspect the pool? Out of scope for v1; add later if needed.
|
||||||
@@ -0,0 +1,43 @@
|
|||||||
|
## Why
|
||||||
|
|
||||||
|
Today every user of `/api/video-agent/*` and `/comfy/*` shares one global `Arc<ComfyUIProcess>` singleton spawned at boot. That means workflows, queue, custom nodes, and uploaded assets are visible across users — unsuitable for the 100+ user multi-tenant deployment we are moving toward. We need per-user isolation now, with an architecture that can move to K3S later without a second rewrite.
|
||||||
|
|
||||||
|
## What Changes
|
||||||
|
|
||||||
|
- **BREAKING (internal API only):** Replace the `state.comfyui: Arc<ComfyUIProcess>` singleton with `state.comfyui_pool: Arc<ComfyUIPool>`. All existing `/api/video-agent/*` handlers refactor to look up (or lazily start) the caller's container via the pool.
|
||||||
|
- Introduce per-user on-demand ComfyUI containers. One container per authenticated user, started on first `/comfy/*` or `/api/video-agent/*` request, reaped after **60 minutes idle**.
|
||||||
|
- Cap the pool at **30 concurrent containers** (≈60 GB RAM at 2 GB each). Requests beyond that queue with a clear "service busy" response and full-screen loading UI hint.
|
||||||
|
- Each container limited to **2 CPU / 2 GB RAM** (`--cpus=2 --memory=2g --memory-swap=2g`), bound to a port from the **8200–8299** pool, named `comfyui-user-<user_id>`.
|
||||||
|
- Persistent per-user volume mounted at `/data/comfyui-users/<user_id>/` so custom workflows, inputs, and outputs survive container restarts.
|
||||||
|
- Add a reverse-proxy route family `/comfy/*` (HTTP + WebSocket) that routes the authenticated user to their own container. **No username in URL** — identity comes from the existing JWT cookie. ComfyUI's HTML/asset paths assume root, so we proxy at the host level rather than under a subpath.
|
||||||
|
- Bake the existing `comfy_source/comfy_api_nodes/util/_helpers.py` env-var fallback patch into the container image so ComfyUI's built-in OpenAI nodes pick up `BOOSAPI_BUILTIN_KEY` automatically.
|
||||||
|
- Abstract the container backend behind a `ContainerRuntime` trait. v1 ships `DockerRuntime`. v2 (when we move to K3S) adds `K3sRuntime` without touching pool, proxy, or route code.
|
||||||
|
- Wire pool startup, idle reaper, and proxy into the existing axum app and `build-and-run.sh`.
|
||||||
|
|
||||||
|
## Capabilities
|
||||||
|
|
||||||
|
### New Capabilities
|
||||||
|
- `comfyui-multi-tenant`: per-user ComfyUI container lifecycle (start, idle reap, max-concurrency, port allocation, persistent volumes).
|
||||||
|
- `container-runtime`: backend-agnostic container abstraction (`ContainerRuntime` trait + `DockerRuntime` impl) so K3S can be added later without rewriting callers.
|
||||||
|
- `comfyui-proxy`: cookie-authenticated reverse proxy at `/comfy/*` (HTTP + WebSocket) that maps the request to the caller's container.
|
||||||
|
- `video-agent-per-user`: refactor of the existing `/api/video-agent/*` handlers to use the per-user pool instead of the global singleton, preserving the current frontend contract.
|
||||||
|
|
||||||
|
### Modified Capabilities
|
||||||
|
<!-- None — no specs/ directory exists yet, so all four capabilities are net-new. -->
|
||||||
|
|
||||||
|
## Impact
|
||||||
|
|
||||||
|
- **Code touched (Rust):**
|
||||||
|
- `src/services/comfyui_process.rs` — keep as the in-container process model used by `DockerRuntime`, no longer spawned directly by `lib.rs`.
|
||||||
|
- `src/services/` — new `comfyui_pool.rs`, `container_runtime.rs`, `docker_runtime.rs`.
|
||||||
|
- `src/routes/video_agent.rs` — every handler switches from `state.comfyui` to `state.comfyui_pool.get_or_start(user_id)`.
|
||||||
|
- `src/routes/` — new `comfy_proxy.rs` mounting `/comfy/*` (HTTP + WS).
|
||||||
|
- `src/lib.rs` — `AppState` field rename, pool construction, idle-reaper task spawn.
|
||||||
|
- **Infra:**
|
||||||
|
- New `docker/Dockerfile.comfyui` baking the `_helpers.py` patch and `bosapi_comfy_nodes.py` into `bosapi-comfyui:multi-tenant`.
|
||||||
|
- `build-and-run.sh` builds the image; the Rust binary itself launches/stops containers at runtime (no docker-compose).
|
||||||
|
- Host paths: `/data/comfyui-users/<user_id>/` (persistent), `/var/run/docker.sock` mounted into the API container if API runs in Docker.
|
||||||
|
- **Auth:** Existing JWT cookie used as-is. New `/comfy/*` proxy reads it; existing `/api/video-agent/*` already does.
|
||||||
|
- **Frontend:** No contract change for `/api/video-agent/*`. New full-screen loading UI shown during the 10–20 s container startup; backend exposes a status endpoint the UI can poll.
|
||||||
|
- **Removals:** Old `docker-compose.comfyui.yml` becomes dev-only; production no longer uses it. The single global `state.comfyui` field disappears.
|
||||||
|
- **Out of scope:** GPU support (we are API-only), per-org quotas (single global cap of 30 only), K3S backend implementation (v2 task).
|
||||||
@@ -0,0 +1,78 @@
|
|||||||
|
## ADDED Requirements
|
||||||
|
|
||||||
|
### Requirement: Per-user container lifecycle
|
||||||
|
The system SHALL maintain at most one ComfyUI container per authenticated user. The container SHALL be started lazily on the user's first request to `/comfy/*` or `/api/video-agent/*` and SHALL be reaped after 60 minutes of inactivity.
|
||||||
|
|
||||||
|
#### Scenario: First request starts the container
|
||||||
|
- **WHEN** an authenticated user with no running container makes a request to `/comfy/*`
|
||||||
|
- **THEN** the system starts a container named `comfyui-user-<user_id>`, allocates a port from 8200–8299, records `last_used_at = now`, and forwards the request once the container's HTTP port is reachable
|
||||||
|
|
||||||
|
#### Scenario: Subsequent request reuses the container
|
||||||
|
- **WHEN** the same user makes another request within 60 minutes
|
||||||
|
- **THEN** the system reuses the existing container, updates `last_used_at = now`, and forwards the request without restarting
|
||||||
|
|
||||||
|
#### Scenario: Idle container is reaped
|
||||||
|
- **WHEN** a container's `last_used_at` is older than 60 minutes AND it has no active WebSocket connections
|
||||||
|
- **THEN** the system stops the container via the runtime's `stop` method and releases its port back to the pool
|
||||||
|
|
||||||
|
### Requirement: Concurrency cap
|
||||||
|
The system SHALL refuse to start a new container when 30 are already running.
|
||||||
|
|
||||||
|
#### Scenario: 31st user is rejected
|
||||||
|
- **WHEN** 30 containers are running and a 31st distinct user requests `/comfy/*`
|
||||||
|
- **THEN** the system responds with HTTP 503 and a `Retry-After: 60` header, and does NOT evict any existing container
|
||||||
|
|
||||||
|
#### Scenario: Cap frees on idle reap
|
||||||
|
- **WHEN** any of the 30 containers is reaped
|
||||||
|
- **THEN** the next pending user request that arrives starts a fresh container in the freed slot
|
||||||
|
|
||||||
|
### Requirement: Resource limits
|
||||||
|
Each container SHALL run with `--cpus=2`, `--memory=2g`, `--memory-swap=2g`, `--pids-limit=512`, `--read-only` rootfs, `--cap-drop=ALL`, and `--security-opt=no-new-privileges`.
|
||||||
|
|
||||||
|
#### Scenario: Container respects CPU and memory caps
|
||||||
|
- **WHEN** a new container is started
|
||||||
|
- **THEN** `docker inspect` shows `HostConfig.NanoCpus = 2000000000`, `HostConfig.Memory = 2147483648`, `HostConfig.MemorySwap = 2147483648`
|
||||||
|
|
||||||
|
#### Scenario: Container has no extra capabilities
|
||||||
|
- **WHEN** a new container is started
|
||||||
|
- **THEN** `docker inspect` shows `HostConfig.CapDrop` contains `ALL` and `HostConfig.SecurityOpt` contains `no-new-privileges`
|
||||||
|
|
||||||
|
### Requirement: Persistent per-user volumes
|
||||||
|
Each container SHALL mount the host directory `/data/comfyui-users/<user_id>/` so that custom workflows, inputs, and outputs survive container restarts.
|
||||||
|
|
||||||
|
#### Scenario: Workflow saved in one session is visible after restart
|
||||||
|
- **WHEN** a user saves a workflow file to ComfyUI's `user/` directory, then triggers a container restart (idle reap + new request)
|
||||||
|
- **THEN** the workflow file is still present in the new container's `user/` directory
|
||||||
|
|
||||||
|
#### Scenario: Volume is per-user
|
||||||
|
- **WHEN** user A saves a file and user B starts their own container
|
||||||
|
- **THEN** user B's container does NOT see user A's file
|
||||||
|
|
||||||
|
### Requirement: Built-in API key injection
|
||||||
|
Each container SHALL receive `BOOSAPI_BUILTIN_KEY`, `BOOSAPI_BASE`, and `BOOSAPI_KEY` as environment variables so ComfyUI's built-in OpenAI nodes (via the patched `_helpers.py`) and the `bosapi_comfy_nodes.py` custom nodes both authenticate successfully.
|
||||||
|
|
||||||
|
#### Scenario: Internal OpenAI node authenticates
|
||||||
|
- **WHEN** a workflow uses a built-in ComfyUI OpenAI node and the container has `BOOSAPI_BUILTIN_KEY` set
|
||||||
|
- **THEN** the node's outbound request to the BoosAPI shim includes `Authorization: Bearer <BOOSAPI_BUILTIN_KEY>`
|
||||||
|
|
||||||
|
#### Scenario: Custom BoosAPI node authenticates
|
||||||
|
- **WHEN** a workflow uses `BoosAPITextGen` from `bosapi_comfy_nodes.py` and the container has `BOOSAPI_KEY` set
|
||||||
|
- **THEN** the node's request to `BOOSAPI_BASE` includes `Authorization: Bearer <BOOSAPI_KEY>`
|
||||||
|
|
||||||
|
### Requirement: Container state observable to callers
|
||||||
|
The pool SHALL expose each container's state as `starting`, `ready`, or `stopped` so the frontend can render a loading UI during cold start.
|
||||||
|
|
||||||
|
#### Scenario: Status reflects starting then ready
|
||||||
|
- **WHEN** a user's first request triggers a container start
|
||||||
|
- **THEN** `pool.state(user_id)` returns `starting` until the HTTP port is reachable, then `ready`
|
||||||
|
|
||||||
|
#### Scenario: Status reflects stopped after reap
|
||||||
|
- **WHEN** a container has been idle-reaped
|
||||||
|
- **THEN** `pool.state(user_id)` returns `stopped`
|
||||||
|
|
||||||
|
### Requirement: Reconcile on startup
|
||||||
|
On API restart, the pool SHALL list all `comfyui-user-*` containers from the runtime and reconstruct its in-memory map (user_id, port, last_used_at).
|
||||||
|
|
||||||
|
#### Scenario: Existing containers are adopted
|
||||||
|
- **WHEN** the API process restarts while a `comfyui-user-42` container is still running
|
||||||
|
- **THEN** the new pool's map contains an entry for user 42 with the original container's port, and the next request from user 42 reuses it
|
||||||
@@ -0,0 +1,45 @@
|
|||||||
|
## ADDED Requirements
|
||||||
|
|
||||||
|
### Requirement: Cookie-authenticated reverse proxy at /comfy/*
|
||||||
|
The system SHALL mount a reverse proxy at `/comfy/*` that resolves the caller's `user_id` from the existing JWT cookie and forwards both HTTP and WebSocket traffic to that user's container.
|
||||||
|
|
||||||
|
#### Scenario: Anonymous request is rejected
|
||||||
|
- **WHEN** a request to `/comfy/` arrives without a valid JWT cookie
|
||||||
|
- **THEN** the proxy responds with HTTP 401 and does NOT touch the container pool
|
||||||
|
|
||||||
|
#### Scenario: Authenticated HTTP request reaches the right container
|
||||||
|
- **WHEN** user 7 requests `GET /comfy/api/object_info`
|
||||||
|
- **THEN** the proxy forwards the request body, query, and headers to `http://127.0.0.1:<port_for_user_7>/api/object_info` and streams the response back unchanged
|
||||||
|
|
||||||
|
#### Scenario: Authenticated WebSocket upgrade is proxied
|
||||||
|
- **WHEN** user 7's browser opens `wss://host/comfy/ws?clientId=abc`
|
||||||
|
- **THEN** the proxy upgrades the connection, opens a WS to `ws://127.0.0.1:<port_for_user_7>/ws?clientId=abc`, and pipes frames in both directions until either side closes
|
||||||
|
|
||||||
|
### Requirement: No URL rewriting of asset paths
|
||||||
|
The proxy SHALL forward request paths to the container with the `/comfy` prefix stripped and SHALL NOT rewrite response bodies.
|
||||||
|
|
||||||
|
#### Scenario: Asset path is preserved
|
||||||
|
- **WHEN** the browser requests `/comfy/extensions/core/foo.js`
|
||||||
|
- **THEN** the proxy forwards `/extensions/core/foo.js` to the container, and the response body is returned byte-for-byte
|
||||||
|
|
||||||
|
#### Scenario: ComfyUI's root-relative URLs work
|
||||||
|
- **WHEN** the container returns HTML referencing `/scripts/app.js`
|
||||||
|
- **THEN** the browser fetches `/comfy/scripts/app.js` (because the page was served under `/comfy/`) and the proxy resolves it correctly
|
||||||
|
|
||||||
|
### Requirement: Cold-start handling
|
||||||
|
If a request arrives while the user's container is `starting`, the proxy SHALL hold the request (up to 30 s) until the container reports `ready`, instead of returning 502.
|
||||||
|
|
||||||
|
#### Scenario: Request waits for cold start
|
||||||
|
- **WHEN** user 7's first request to `/comfy/` arrives and the container needs 12 s to start
|
||||||
|
- **THEN** the proxy waits for the readiness signal and forwards the request once the container is ready, returning a normal 200
|
||||||
|
|
||||||
|
#### Scenario: Cold-start timeout
|
||||||
|
- **WHEN** a container fails to become ready within 30 s
|
||||||
|
- **THEN** the proxy responds with HTTP 504 and the pool marks the container `failed` for the next idle scan
|
||||||
|
|
||||||
|
### Requirement: Last-used timestamp updated on every frame
|
||||||
|
Every HTTP request and WebSocket frame proxied to the container SHALL update `pool.containers[user_id].last_used_at = now`.
|
||||||
|
|
||||||
|
#### Scenario: Long-running WS keeps container alive
|
||||||
|
- **WHEN** a user holds a WS connection open for 90 minutes while exchanging frames every minute
|
||||||
|
- **THEN** the container is not reaped at the 60-minute mark
|
||||||
@@ -0,0 +1,34 @@
|
|||||||
|
## ADDED Requirements
|
||||||
|
|
||||||
|
### Requirement: Backend-agnostic ContainerRuntime trait
|
||||||
|
The system SHALL expose container lifecycle behind a single `ContainerRuntime` trait so the pool, proxy, and route handlers do NOT depend on Docker directly.
|
||||||
|
|
||||||
|
#### Scenario: Pool uses only the trait
|
||||||
|
- **WHEN** the source for `services/comfyui_pool.rs` is inspected
|
||||||
|
- **THEN** it imports `ContainerRuntime` and `ContainerSpec` only — never `bollard`, `docker`, or `kube` types
|
||||||
|
|
||||||
|
#### Scenario: Trait surface is sufficient for v1
|
||||||
|
- **WHEN** a `DockerRuntime` is provided
|
||||||
|
- **THEN** the pool can start containers, stop them, check liveness, and list existing `comfyui-user-*` containers using only the trait's methods
|
||||||
|
|
||||||
|
### Requirement: DockerRuntime implementation
|
||||||
|
The system SHALL ship a `DockerRuntime` that implements `ContainerRuntime` against the local Docker daemon, either via the host socket `/var/run/docker.sock` or the `docker` CLI.
|
||||||
|
|
||||||
|
#### Scenario: Start spawns a container
|
||||||
|
- **WHEN** `DockerRuntime::start(spec)` is called with `image=bosapi-comfyui:multi-tenant`, `name=comfyui-user-7`, port mapping, env, and volume mounts
|
||||||
|
- **THEN** a container with that name exists, is running, and exposes the chosen host port
|
||||||
|
|
||||||
|
#### Scenario: Stop removes the container
|
||||||
|
- **WHEN** `DockerRuntime::stop("comfyui-user-7")` is called
|
||||||
|
- **THEN** the container is stopped and removed (image used `--rm`-equivalent semantics) and its name becomes reusable
|
||||||
|
|
||||||
|
#### Scenario: List finds existing user containers
|
||||||
|
- **WHEN** `DockerRuntime::list_user_containers()` is called
|
||||||
|
- **THEN** it returns one entry per running container whose name starts with `comfyui-user-`, including its port mapping
|
||||||
|
|
||||||
|
### Requirement: K3S extension point reserved
|
||||||
|
The trait SHALL be designed such that a future `K3sRuntime` implementing the same surface can replace `DockerRuntime` without changes to the pool, proxy, or route handlers.
|
||||||
|
|
||||||
|
#### Scenario: Swapping runtimes does not break callers
|
||||||
|
- **WHEN** `AppState::comfyui_pool` is constructed with a `Box<dyn ContainerRuntime>` instead of a concrete `DockerRuntime`
|
||||||
|
- **THEN** all callers compile and behavior is unchanged for the `DockerRuntime` instance
|
||||||
@@ -0,0 +1,41 @@
|
|||||||
|
## ADDED Requirements
|
||||||
|
|
||||||
|
### Requirement: video-agent handlers route through the pool
|
||||||
|
Every handler under `routes/video_agent.rs` SHALL resolve the caller's `user_id` from the JWT cookie and operate on `state.comfyui_pool.get_or_start(user_id)` instead of the global `state.comfyui` singleton. The wire format of existing endpoints SHALL NOT change.
|
||||||
|
|
||||||
|
#### Scenario: status reflects per-user container
|
||||||
|
- **WHEN** user 7 calls `GET /api/video-agent/status`
|
||||||
|
- **THEN** the response describes user 7's container only, including a new `containerState: "starting" | "ready" | "stopped"` field, and the pre-existing JSON keys are preserved
|
||||||
|
|
||||||
|
#### Scenario: start is idempotent per user
|
||||||
|
- **WHEN** user 7 calls `POST /api/video-agent/start` twice in a row
|
||||||
|
- **THEN** the first call starts (or reuses) the container and returns 200, the second call returns 200 immediately without spawning a duplicate container
|
||||||
|
|
||||||
|
#### Scenario: stop only stops the caller's container
|
||||||
|
- **WHEN** user 7 calls `POST /api/video-agent/stop`
|
||||||
|
- **THEN** user 7's container is stopped and removed, and any other user's running container is unaffected
|
||||||
|
|
||||||
|
#### Scenario: chat / analyze / generate-* call the caller's container
|
||||||
|
- **WHEN** user 7 calls any of `/api/video-agent/{chat,analyze,character-prompts,generate-character,scene-prompt,generate-scene,generate-tts,full-workflow}`
|
||||||
|
- **THEN** the handler forwards the request to user 7's container's HTTP endpoint and returns the response
|
||||||
|
|
||||||
|
### Requirement: Backwards-compatible response shape
|
||||||
|
Existing JSON keys returned by `/api/video-agent/*` SHALL remain present and identical to the singleton-era response. Only additive fields are allowed.
|
||||||
|
|
||||||
|
#### Scenario: status payload is a superset
|
||||||
|
- **WHEN** the frontend (unchanged) parses the response of `/api/video-agent/status`
|
||||||
|
- **THEN** every field it currently reads is still present with the same type and meaning; only `containerState` (and any new optional fields) are added
|
||||||
|
|
||||||
|
### Requirement: Authentication required
|
||||||
|
All `/api/video-agent/*` routes SHALL require a valid JWT cookie. Anonymous requests SHALL return HTTP 401 without touching the pool.
|
||||||
|
|
||||||
|
#### Scenario: Anonymous request blocked
|
||||||
|
- **WHEN** an unauthenticated client calls `POST /api/video-agent/start`
|
||||||
|
- **THEN** the response is HTTP 401 and no container is started
|
||||||
|
|
||||||
|
### Requirement: Removal of global state.comfyui
|
||||||
|
The `comfyui: Arc<ComfyUIProcess>` field on `AppState` SHALL be removed once `comfyui_pool` is the sole code path. No call site SHALL retain a reference to the singleton.
|
||||||
|
|
||||||
|
#### Scenario: Codebase contains no singleton refs
|
||||||
|
- **WHEN** the source tree is grepped for `state.comfyui` (without the `_pool` suffix)
|
||||||
|
- **THEN** there are zero matches outside of historical comments
|
||||||
@@ -0,0 +1,69 @@
|
|||||||
|
## 1. ContainerRuntime trait + DockerRuntime
|
||||||
|
|
||||||
|
- [x] 1.1 Add `services/container_runtime.rs` with `ContainerRuntime` trait, `ContainerSpec` (image, name, env, port_map, mounts, cpu, mem, mem_swap, pids_limit, read_only, cap_drop, security_opt), `ContainerHandle` (id, host_port).
|
||||||
|
- [x] 1.2 Add `services/docker_runtime.rs` implementing `ContainerRuntime` via `bollard` (or `tokio::process::Command` shelling to `docker` if `bollard` is heavy) — `start`, `stop`, `is_alive`, `list_user_containers`.
|
||||||
|
- [ ] 1.3 Unit-test `DockerRuntime` against the local Docker daemon: start/stop/list cycle on a hello-world image; assert resource flags via `inspect`. _(skipped this session — needs a live Docker daemon)_
|
||||||
|
|
||||||
|
## 2. ComfyUIPool
|
||||||
|
|
||||||
|
- [x] 2.1 Add `services/comfyui_pool.rs` with `ComfyUIPool { runtime: Arc<dyn ContainerRuntime>, containers: RwLock<HashMap<i64, PoolEntry>>, ports: PortAllocator(8200..=8299), max: usize, idle_timeout: Duration }`.
|
||||||
|
- [x] 2.2 Implement `get_or_start(user_id) -> Result<ContainerHandle>`: dedupe in-flight starts via per-user `Mutex<Option<oneshot::Receiver>>`; allocate port; build `ContainerSpec` with `image=bosapi-comfyui:multi-tenant`, name `comfyui-user-<id>`, env (`BOOSAPI_BUILTIN_KEY`, `BOOSAPI_BASE`, `BOOSAPI_KEY`), volumes for `/data/comfyui-users/<id>/{input,output,temp,user,custom_nodes_user}`, resource caps from §design Decision 5.
|
||||||
|
- [x] 2.3 Readiness probe: poll `GET http://127.0.0.1:<port>/system_stats` until 200 or 30 s timeout; mark state `ready`/`failed`.
|
||||||
|
- [x] 2.4 Implement `stop(user_id)`, `state(user_id) -> "starting"|"ready"|"stopped"|"failed"`, `touch(user_id)`.
|
||||||
|
- [x] 2.5 Enforce concurrency cap: when `containers.len() >= 30`, return `PoolError::AtCapacity`; map to HTTP 503 `Retry-After: 60` at the proxy/route layer.
|
||||||
|
- [x] 2.6 Spawn idle-reaper tokio task: every 60 s, `stop` containers whose `last_used_at` > 60 min AND have no active WS connections.
|
||||||
|
- [x] 2.7 Reconcile on startup: `runtime.list_user_containers()` → seed `containers` map with existing `comfyui-user-*` entries.
|
||||||
|
|
||||||
|
## 3. AppState wiring
|
||||||
|
|
||||||
|
- [x] 3.1 Add `comfyui_pool: Arc<ComfyUIPool>` to `AppState` in `src/lib.rs` (keep `comfyui` field temporarily behind `cfg(feature = "legacy_comfy_singleton")` for rollback).
|
||||||
|
- [x] 3.2 Construct the pool in `lib.rs::build_app_state` from `DockerRuntime::new()` and the existing `BuiltinKeys.comfyui_agent` secret.
|
||||||
|
- [x] 3.3 Spawn the idle-reaper task at startup; add a graceful-shutdown hook that `stop`s any container we started in this process.
|
||||||
|
|
||||||
|
## 4. /comfy/* reverse proxy
|
||||||
|
|
||||||
|
- [x] 4.1 Add `routes/comfy_proxy.rs` mounting `any("/comfy/*path")` and `get("/comfy/ws")` (WS upgrade) using `axum::extract::ws` + `reqwest` (HTTP) / `tokio_tungstenite` (WS-to-WS).
|
||||||
|
- [x] 4.2 Auth middleware: extract JWT cookie, resolve `user_id`; on missing/invalid, return 401 without touching the pool.
|
||||||
|
- [x] 4.3 HTTP forwarder: strip `/comfy` prefix; copy method, path, query, headers (drop hop-by-hop), body; stream response back unchanged.
|
||||||
|
- [x] 4.4 WebSocket forwarder: open server-side WS to `ws://127.0.0.1:<port>/ws?...`; pump frames bidirectionally; close on either-side close.
|
||||||
|
- [x] 4.5 Cold-start hold: if `pool.state(user_id)` is `starting`, await readiness up to 30 s before forwarding; on timeout return 504.
|
||||||
|
- [x] 4.6 Touch `pool` on every HTTP request and every inbound/outbound WS frame.
|
||||||
|
- [x] 4.7 Register the router in `lib.rs::build_router` so `/comfy/*` is reachable.
|
||||||
|
|
||||||
|
## 5. /api/video-agent/* refactor
|
||||||
|
|
||||||
|
- [x] 5.1 In `routes/video_agent.rs`, add a helper `resolve_user_id(headers) -> Result<i64, AppError>` reading the JWT cookie.
|
||||||
|
- [x] 5.2 Rewrite `status` to call `pool.state(user_id)` and (if ready) the container's HTTP endpoint; preserve every existing JSON key, add `containerState`.
|
||||||
|
- [x] 5.3 Rewrite `start_comfyui` to `pool.get_or_start(user_id)` (idempotent).
|
||||||
|
- [x] 5.4 Rewrite `stop_comfyui` and `restart_comfyui` to act on the caller's container only.
|
||||||
|
- [x] 5.5 Rewrite `chat`, `analyze`, `character_prompts`, `generate_character`, `scene_prompt`, `generate_scene`, `generate_tts`, `full_workflow` to forward to `http://127.0.0.1:<port>/...` for the caller's container.
|
||||||
|
- [x] 5.6 Anonymous-request guard: every handler returns 401 if `resolve_user_id` fails.
|
||||||
|
|
||||||
|
## 6. Image build
|
||||||
|
|
||||||
|
- [x] 6.1 Add `docker/Dockerfile.comfyui` based on user's existing `bosapi-comfyui:ready` lineage; copy patched `comfy_source/comfy_api_nodes/util/_helpers.py` and `bosapi_comfy_nodes.py` into `/workspace/ComfyUI/`. _(implemented as `docker/Dockerfile.comfyui.multi-tenant` to keep the legacy singleton image untouched until rollout 9.4)_
|
||||||
|
- [x] 6.2 ENTRYPOINT script reads `BOOSAPI_BUILTIN_KEY`/`BOOSAPI_BASE`/`BOOSAPI_KEY` and execs `python main.py --listen 0.0.0.0 --port 8188 --comfy-api-base "$BOOSAPI_BASE/api/comfyui-agent"`. _(`docker/comfyui-entrypoint.sh`)_
|
||||||
|
- [x] 6.3 Tag image as `bosapi-comfyui:multi-tenant`; update `build-and-run.sh` to build it before launching the API.
|
||||||
|
- [ ] 6.4 Smoke-test: `docker run --rm --cpus=2 --memory=2g -e BOOSAPI_BUILTIN_KEY=... -p 18188:8188 bosapi-comfyui:multi-tenant` and curl `/system_stats`.
|
||||||
|
|
||||||
|
## 7. Frontend integration
|
||||||
|
|
||||||
|
- [ ] 7.1 Update video-agent panel to render full-screen loading overlay when `containerState === "starting"`; poll `/api/video-agent/status` every 2 s.
|
||||||
|
- [ ] 7.2 Add a `/comfy` entry point (link or iframe) gated on the user being authenticated; show `containerState` while loading.
|
||||||
|
- [ ] 7.3 Handle 503 from `/comfy/*` and `/api/video-agent/*` by showing a "服务繁忙，请稍候" page and retrying after `Retry-After`.
|
||||||
|
|
||||||
|
## 8. Verification
|
||||||
|
|
||||||
|
- [x] 8.1 `cargo check --all-targets` green.
|
||||||
|
- [ ] 8.2 `cargo test` green for new pool / proxy / route tests; pre-existing failures from task #27 explicitly out of scope.
|
||||||
|
- [ ] 8.3 Manual two-user test on staging: user A uploads image into ComfyUI, user B starts a fresh container and confirms A's file is NOT visible.
|
||||||
|
- [ ] 8.4 Manual cold-start test: stop user A's container, hit `/comfy/`, confirm full-screen loader shows for ~10–20 s, then ComfyUI loads.
|
||||||
|
- [ ] 8.5 Manual cap test: start 30 containers, confirm 31st request returns 503 with `Retry-After`.
|
||||||
|
- [ ] 8.6 Idle-reap test: leave a container idle for 65 minutes, confirm `docker ps` no longer lists it and the next request cold-starts.
|
||||||
|
|
||||||
|
## 9. Rollout
|
||||||
|
|
||||||
|
- [x] 9.1 Land 1–6 behind `multi_tenant_comfy` feature flag (default off) so legacy singleton path keeps working.
|
||||||
|
- [ ] 9.2 Enable in staging; run §8 manual tests.
|
||||||
|
- [ ] 9.3 Enable in production; monitor pool metrics for 1 week.
|
||||||
|
- [ ] 9.4 Remove `state.comfyui` singleton field, the feature flag, and any compat branches; archive legacy `docker-compose.comfyui.yml` references.
|
||||||
@@ -0,0 +1,2 @@
|
|||||||
|
schema: spec-driven
|
||||||
|
created: 2026-05-19
|
||||||
@@ -0,0 +1,81 @@
|
|||||||
|
## Context
|
||||||
|
|
||||||
|
`src/web/App.tsx` 当前 1135 行，其中第 132–500+ 行的 `Login` 组件混合了三种登录方式（管理员 token、管理员账号密码、普通用户账号密码）和品牌侧栏文案；`src/web/pages/UserLogin.tsx` 又是同一类登录的重复实现。错误文案分别在 `loginError.ts`（仅 token 登录）、`App.tsx` 内联（账号密码登录）、`UserLogin.tsx` 内联三处解析；后端 `routes/auth.rs` 同时挂在 `/api/auth/login` 与 `/api/users/login` 两个路径上输出同一 JSON 形状。
|
||||||
|
|
||||||
|
OAuth 现状：`src/routes/oauth.rs` 顶部注释明确 "Simplified implementation — complex OAuth flow operations return 501"。`start_provider_flow` / `get_session` / `manual_callback` 都直接返回 "not implemented"，没有 session 存储数据结构（`rg "oauth_sessions"` 零命中）。把 OAuth 当登录入口需要先把整个 provider flow 从 stub 写到可用，与本次前端拆分不在同一量级，因此从本次范围拆出。
|
||||||
|
|
||||||
|
约束：保留现栈（React 18 + 自研 `authSession.ts`），不引入路由库（已在用 `react-router-dom`）外的状态库；不改造桌面端 deeplink。
|
||||||
|
|
||||||
|
## Goals / Non-Goals
|
||||||
|
|
||||||
|
**Goals:**
|
||||||
|
- 把 `App.tsx` 内的登录组件拆出，得到 `pages/auth/LoginShell` + `AdminLoginPage` + `UserLoginPage` 三个文件，单文件不超过 ~250 行。
|
||||||
|
- 所有登录入口（管理员 token、管理员/用户账号密码）的失败提示走同一个 `resolveLoginErrorMessage`，状态码 → 文案映射集中在一处。
|
||||||
|
- `authSession.ts` 暴露 `commitSession({ audience, token })` / `clearSession(audience)` / `getSessionToken(audience)` 入口，所有登录路径只调用它，旧 API 标 `@deprecated` 但保留薄壳。
|
||||||
|
- `UserLoginPage` 底部留 `<OAuthSlot />` 占位组件，本次只渲染容器不挂行为，为下一个 change 留出 UI 落点。
|
||||||
|
|
||||||
|
**Non-Goals:**
|
||||||
|
- 不实现任何 OAuth provider flow（前后端皆不动）。
|
||||||
|
- 不替换 React、不引入 zustand/jotai 或 react-hook-form。
|
||||||
|
- 不重写 `routes/oauth.rs`。
|
||||||
|
- 不动桌面端（electron）登录链路。
|
||||||
|
- 不做 MFA / 邮箱验证码 / 注册流程改造（注册页不在本次范围）。
|
||||||
|
|
||||||
|
## Decisions
|
||||||
|
|
||||||
|
### 1. 拆分粒度：抽 `LoginShell` + 两个 page 组件，而不是一个统一 Login
|
||||||
|
|
||||||
|
`LoginShell` 负责品牌侧栏、容器、i18n 注入、错误展示槽位；`AdminLoginPage` 持有 token/账密两种 mode 切换；`UserLoginPage` 只有账密一种 mode + OAuth 槽位（占位）。
|
||||||
|
|
||||||
|
**为何不做单一 `<Login audience>` 组件**：当前实现就是这种思路，却让 props 组合（audience × mode × OAuth）爆炸。两个 page 组件之间共享 30% 代码（表单原语 + 提交按钮），用 `LoginShell` + 表单子组件复用即可，不必塞回一个组件。
|
||||||
|
|
||||||
|
**为何不直接删 `pages/UserLogin.tsx`**：保留路由 `/user/login` 让用户书签可用，但内部薄壳化，挂 `UserLoginPage` 即可。
|
||||||
|
|
||||||
|
### 2. 登录错误统一：扩展 `loginError.ts`，对 axios-style 异常归一
|
||||||
|
|
||||||
|
现状 `resolveLoginErrorMessage(status, reason)` 只覆盖 401/403/5xx 三类。新增：
|
||||||
|
- 400：参数错误，区分"邮箱格式不合法"/"密码不能为空"两种 reason（后端已有 `bad_request` 文案）。
|
||||||
|
- 409：账号被锁定 / 多次失败（后端目前没有，留接口以便后续接入）。
|
||||||
|
- 网络错误（fetch reject）：返回 `网络异常，请检查连接`。
|
||||||
|
|
||||||
|
`AdminLogin` 与 `UserLogin` 都通过 `await loginRequest(...)` 拿统一形状 `{ status, reason }`，再交给 `resolveLoginErrorMessage`。
|
||||||
|
|
||||||
|
**为何不引入新错误库**：错误形状只有 5 种分支，多一层抽象反而模糊。
|
||||||
|
|
||||||
|
### 3. 会话写入收敛到 `authSession.commitSession`
|
||||||
|
|
||||||
|
新签名：
|
||||||
|
```ts
|
||||||
|
type Audience = "admin" | "user";
|
||||||
|
function commitSession(audience: Audience, token: string, ttlMs?: number, storage?, now?): void;
|
||||||
|
function clearSession(audience: Audience, storage?): void;
|
||||||
|
function getSessionToken(audience: Audience, storage?, now?): string | null;
|
||||||
|
```
|
||||||
|
|
||||||
|
底层仍然用现有 `AUTH_TOKEN_*` / `USER_JWT_*` 两套 key，对外暴露统一 API。`persistAuthSession` / `persistUserSession` 标记为 `@deprecated`，本次先不删（避免影响其他 page 直接调用的地方），下一次清理。
|
||||||
|
|
||||||
|
### 4. OAuth 占位：组件级 `OAuthSlot`，目前只渲染容器
|
||||||
|
|
||||||
|
`UserLoginPage` 底部嵌一个 `<OAuthSlot variant="user" />`，组件内部本次只输出一个空 `<section>`（或带 i18n 文案的 placeholder），不调任何 OAuth API。后续 change 在同一组件里实现 provider 列表 + 跳转 + exchange 链路。
|
||||||
|
|
||||||
|
**为何不直接留空白**：留可识别的 DOM 节点能让后续 change 用 `data-testid="oauth-slot"` 精准插入；留 `OAuthSlot` 文件还能让 `UserLoginPage` 的 import 形态稳定，避免后续 change 又改 page 组件 import。
|
||||||
|
|
||||||
|
## Risks / Trade-offs
|
||||||
|
|
||||||
|
- **[Risk] 拆分过程中 i18n key 错位**：现有 `t('登录')` 等中文 key 和 `i18n.tsx` supplement 强耦合，移动文件时容易丢翻译。→ Mitigation：拆分时一次性把 `t('xxx')` 调用搬到新文件，并跑 `i18n.test.ts` 校验。
|
||||||
|
- **[Risk] 测试迁移漏掉断言**：`App.login-surface.test.tsx` 是基于 `App.tsx` 内的 DOM 断言，组件路径变了断言要重写。→ Mitigation：新建 `pages/auth/LoginShell.test.tsx` 把原断言搬过去，老测试文件保留薄壳验证 `App.tsx` 里挂载入口仍存在。
|
||||||
|
- **[Trade-off] 保留 `persistAuthSession` / `persistUserSession`**：增加一次过渡期的 API 重复，但避免本次 PR 触碰过多调用点。
|
||||||
|
- **[Trade-off] OAuth 占位组件本次不挂任何行为**：会有一个看起来空的 DOM 区域，但避免本次 PR 误把"占位"当"已接入"看待；文档段落明确说明。
|
||||||
|
|
||||||
|
## Migration Plan
|
||||||
|
|
||||||
|
1. **第 1 步（无破坏）**：扩展 `loginError.ts` + 新建 `authSession.commitSession`，老函数转为薄壳调用新函数。跑全量单测确认无回归。
|
||||||
|
2. **第 2 步（前端拆分）**：新建 `src/web/pages/auth/LoginShell.tsx` + `AdminLoginPage.tsx` + `UserLoginPage.tsx` + `OAuthSlot.tsx`（占位）；`App.tsx` 改为挂 `AdminLoginPage`；`pages/UserLogin.tsx` 内部 import `UserLoginPage`。手测三种现有登录路径。
|
||||||
|
3. **第 3 步（清理）**：删除 `App.tsx` 中已无用的 Login 组件 dead code，更新 `App.login-surface.test.tsx`，更新文档。
|
||||||
|
|
||||||
|
回滚：每步独立提交；如果第 2 步出问题，revert 即可恢复 `App.tsx` 内联实现。
|
||||||
|
|
||||||
|
## Open Questions
|
||||||
|
|
||||||
|
- 桌面端（electron）登录是否也通过同一组件？本次明确不在范围，但 `LoginShell` 不能假设 `window.location.origin` 是 web，组件内不出现 `window.location` 直接访问。
|
||||||
|
- `OAuthSlot` 是否需要通过 feature flag 控制可见性？本次按"始终渲染容器、内部留 placeholder"处理，下一个 change 再决定。
|
||||||
@@ -0,0 +1,34 @@
|
|||||||
|
## Why
|
||||||
|
|
||||||
|
`src/web/App.tsx` 已经膨胀到 1135 行，其中 `Login` 组件（约 132–500+ 行）耦合了管理员 token 登录、管理员账号密码登录、用户账号密码登录、品牌侧栏、能力位文案和 i18n 适配，与同文件里另一个用户专用的 `pages/UserLogin.tsx` 形成两套并行实现。错误提示分散在三处（`loginError.ts`、`App.tsx` 内联解析、`UserLogin.tsx` 内联解析），状态码 / 文案不一致。
|
||||||
|
|
||||||
|
继续在 `App.tsx` 内堆登录逻辑会让后续接入新的登录方式（OAuth 第三方登录、邮箱验证码等）越来越难，也会让"登录失败"这类用户能看到的错误持续不一致。
|
||||||
|
|
||||||
|
OAuth 第三方登录入口本来与本次重构同期推进，但实施前确认 `src/routes/oauth.rs` 全套 provider flow 仍是 stub（所有 handler 返回 "not implemented"），缺少 session 存储、callback 解析、token 兑换链路。把 OAuth 后端从 stub 写到可用的体量与本次拆分不成比例，因此从本次范围中拆出，单独立项。
|
||||||
|
|
||||||
|
## What Changes
|
||||||
|
|
||||||
|
- 把 `App.tsx` 内的 `Login` 组件拆分成独立的 `pages/auth/` 模块（`AdminLoginPage` / `UserLoginPage` 复用同一套 `LoginShell` + 表单原语），合并 `pages/UserLogin.tsx` 的重复实现。
|
||||||
|
- 新增一个统一的前端登录错误解析层（扩展 `loginError.ts`），覆盖 `/api/auth/login`、`/api/users/login` 两个端点的所有失败码，所有登录入口走同一函数获得文案。
|
||||||
|
- 统一 `authSession.ts` 中 admin/user 两套 token 的写入路径，登录成功后通过 `commitSession(audience, token)` 单点持久化，避免页面各自调用 `persistAuthSession` / `persistUserSession`。
|
||||||
|
- 在用户登录页底部预留一个 OAuth 占位区域（不挂任何 provider 按钮），为后续接入 OAuth signin 留出 UI 槽位。
|
||||||
|
|
||||||
|
不在本次范围：
|
||||||
|
- OAuth provider flow 的实际接入（前后端均不实现），由独立 change 跟进。
|
||||||
|
- 替换 React、引入新状态库、改造桌面端 deeplink、邮箱验证码、注册流程改造。
|
||||||
|
|
||||||
|
## Capabilities
|
||||||
|
|
||||||
|
### New Capabilities
|
||||||
|
- `frontend-login`: 前端登录界面与会话建立流程（管理员 token / 账号密码、错误展示、会话持久化、OAuth 占位）。
|
||||||
|
|
||||||
|
### Modified Capabilities
|
||||||
|
<!-- 当前 openspec/specs/ 目录为空，没有既有 spec 受影响。 -->
|
||||||
|
|
||||||
|
## Impact
|
||||||
|
|
||||||
|
- 前端：`src/web/App.tsx`（拆出 Login）、`src/web/pages/UserLogin.tsx`（薄壳化）、`src/web/loginError.ts`（扩展）、`src/web/authSession.ts`（收敛入口）、新增 `src/web/pages/auth/`。
|
||||||
|
- 后端：本次不动 `routes/auth.rs` / `routes/oauth.rs`。
|
||||||
|
- 数据库：本次无迁移。
|
||||||
|
- 测试：`App.login-surface.test.tsx`、`loginError.test.ts` 需迁移到新组件位置。
|
||||||
|
- 文档：`README.md` / `AGENTS.md` 中"登录方式"段落需更新（标注 OAuth 由后续 change 接入）。
|
||||||
@@ -0,0 +1,77 @@
|
|||||||
|
## ADDED Requirements
|
||||||
|
|
||||||
|
### Requirement: 登录界面拆分为独立模块
|
||||||
|
|
||||||
|
前端登录界面 SHALL 由 `src/web/pages/auth/` 目录下的独立模块承担，`App.tsx` 不再内联实现登录表单。`pages/auth/` 至少包含 `LoginShell`（品牌侧栏 + 容器 + 错误槽位）、`AdminLoginPage`（管理员 token / 账密两种 mode）、`UserLoginPage`（用户账密 + OAuth 占位槽）三个组件，单文件不超过 250 行。`pages/UserLogin.tsx` MUST 退化为薄壳，仅渲染 `UserLoginPage`，以保留既有路由 `/user/login`。
|
||||||
|
|
||||||
|
#### Scenario: 管理员入口挂载新组件
|
||||||
|
|
||||||
|
- **WHEN** 用户访问根路由 `/` 且未登录
|
||||||
|
- **THEN** 应用渲染 `AdminLoginPage`，且页面来源文件位于 `src/web/pages/auth/AdminLoginPage.tsx`
|
||||||
|
|
||||||
|
#### Scenario: 用户入口保留旧路由
|
||||||
|
|
||||||
|
- **WHEN** 用户访问 `/user/login`
|
||||||
|
- **THEN** 应用渲染 `UserLoginPage`，且 `src/web/pages/UserLogin.tsx` 内部仅作 import + 转发，不再持有表单状态
|
||||||
|
|
||||||
|
#### Scenario: App.tsx 不再持有登录表单状态
|
||||||
|
|
||||||
|
- **WHEN** 在 `src/web/App.tsx` 中查找形如 `setEmail` / `setPassword` / `setToken` 的 useState
|
||||||
|
- **THEN** 这些 state hooks 不出现在 `App.tsx` 中（已迁移到新组件）
|
||||||
|
|
||||||
|
### Requirement: 统一登录错误解析
|
||||||
|
|
||||||
|
所有登录入口（管理员 token、管理员账密、用户账密）的失败提示 SHALL 通过 `src/web/loginError.ts` 中的 `resolveLoginErrorMessage(status, reason)` 解析。该函数 MUST 覆盖 400 / 401 / 403 / 409 / 5xx 五类状态，以及 `network` 这类网络异常 sentinel。组件层 MUST NOT 自行 `switch (status)` 写文案。
|
||||||
|
|
||||||
|
#### Scenario: 401 返回统一令牌无效文案
|
||||||
|
|
||||||
|
- **WHEN** 调用 `resolveLoginErrorMessage(401, "")`
|
||||||
|
- **THEN** 返回包含"令牌无效"或"账号或密码错误"语义的中文文案
|
||||||
|
|
||||||
|
#### Scenario: 网络异常返回统一文案
|
||||||
|
|
||||||
|
- **WHEN** 前端 fetch 抛错（无 status 可读），调用 `resolveLoginErrorMessage("network", "")`
|
||||||
|
- **THEN** 返回"网络异常，请检查连接"文案
|
||||||
|
|
||||||
|
#### Scenario: 400 携带 reason 时分流文案
|
||||||
|
|
||||||
|
- **WHEN** 调用 `resolveLoginErrorMessage(400, "invalid email format")`
|
||||||
|
- **THEN** 返回"邮箱格式不合法"文案；当 reason 为空时返回通用"请求参数错误"
|
||||||
|
|
||||||
|
### Requirement: 会话写入收敛为单一 API
|
||||||
|
|
||||||
|
`src/web/authSession.ts` SHALL 暴露 `commitSession(audience, token, ttlMs?, storage?, now?)` 与 `clearSession(audience, storage?)` 与 `getSessionToken(audience, storage?, now?)` 三个函数。所有登录页 MUST 仅通过这三个函数读写本地存储。`persistAuthSession` 与 `persistUserSession` 仍保留为 `@deprecated` 薄壳调用 `commitSession`。
|
||||||
|
|
||||||
|
#### Scenario: 管理员登录成功调用 commitSession
|
||||||
|
|
||||||
|
- **WHEN** 管理员账密登录成功，组件拿到 JWT
|
||||||
|
- **THEN** 组件调用 `commitSession("admin", jwt)`，且不直接调用 `localStorage.setItem`
|
||||||
|
|
||||||
|
#### Scenario: 用户登录成功调用 commitSession
|
||||||
|
|
||||||
|
- **WHEN** 用户账密登录成功，组件拿到 JWT
|
||||||
|
- **THEN** 组件调用 `commitSession("user", jwt)`，且不直接调用 `persistUserSession`
|
||||||
|
|
||||||
|
#### Scenario: 旧 API 仍可工作
|
||||||
|
|
||||||
|
- **WHEN** 现有调用方仍调用 `persistAuthSession(localStorage, token)`
|
||||||
|
- **THEN** token 仍按原路径写入 `localStorage`，行为与重构前一致
|
||||||
|
|
||||||
|
### Requirement: 登录页预留 OAuth 占位槽
|
||||||
|
|
||||||
|
`UserLoginPage` SHALL 在表单下方渲染一个名为 `OAuthSlot` 的占位组件作为后续 OAuth 登录入口的挂载点。本次实施 `OAuthSlot` MUST NOT 调用任何 OAuth 相关 API，也 MUST NOT 触发跳转；它只渲染一个带稳定 `data-testid="oauth-slot"` 的容器节点（可包含 i18n 文案 placeholder）。`AdminLoginPage` MUST NOT 包含该槽位。
|
||||||
|
|
||||||
|
#### Scenario: 用户登录页渲染占位槽
|
||||||
|
|
||||||
|
- **WHEN** 用户访问 `/user/login` 加载 `UserLoginPage`
|
||||||
|
- **THEN** DOM 中存在 `data-testid="oauth-slot"` 节点，且不发起任何 `/api/oauth/*` 请求
|
||||||
|
|
||||||
|
#### Scenario: 管理员登录页不含占位槽
|
||||||
|
|
||||||
|
- **WHEN** 用户访问根路由 `/` 加载 `AdminLoginPage`
|
||||||
|
- **THEN** DOM 中不存在 `data-testid="oauth-slot"` 节点
|
||||||
|
|
||||||
|
#### Scenario: 占位槽不触发任何行为
|
||||||
|
|
||||||
|
- **WHEN** 用户与 `OAuthSlot` 内任何子元素交互
|
||||||
|
- **THEN** 不发出网络请求，不修改 `authSession`，不触发路由跳转
|
||||||
@@ -0,0 +1,48 @@
|
|||||||
|
## 1. 基础工具层（无破坏改造）
|
||||||
|
|
||||||
|
- [x] 1.1 在 `src/web/loginError.ts` 中扩展 `resolveLoginErrorMessage`，覆盖 400（区分 `invalid email format` / 通用）、409（`no_account`）、`"network"` sentinel；保留现有 401/403/5xx 行为
|
||||||
|
- [x] 1.2 更新 `src/web/loginError.test.ts`，为新增分支补单测（400+reason、409、network、5xx 已有保留）
|
||||||
|
- [x] 1.3 在 `src/web/authSession.ts` 新增 `commitSession(audience, token, ttlMs?, storage?, now?)` / `clearSession(audience, storage?)` / `getSessionToken(audience, storage?, now?)`，内部委托给现有 `AUTH_TOKEN_*` / `USER_JWT_*` 路径
|
||||||
|
- [x] 1.4 把 `persistAuthSession` / `persistUserSession` / `clearAuthSession` / `clearUserSession` 改为薄壳调用 `commitSession` / `clearSession`，加 `@deprecated` JSDoc
|
||||||
|
- [x] 1.5 扩充 `src/web/authSession.test.ts`，覆盖 `commitSession("admin"|"user")` 双路径与旧 API 仍可工作
|
||||||
|
|
||||||
|
## 2. 前端组件拆分
|
||||||
|
|
||||||
|
- [x] 2.1 新建目录 `src/web/pages/auth/`，加 `LoginShell.tsx`：抽出品牌侧栏（logo / 文案 / capability list / 页脚链接）+ 容器 + 错误展示槽位，接收 `children` + `error` prop
|
||||||
|
- [x] 2.2 新建 `src/web/pages/auth/AdminLoginPage.tsx`，承接现有 `App.tsx` 中的 `Login`（admin 分支：token + 账密两种 mode 切换），所有错误走 `resolveLoginErrorMessage`，写入会话调 `commitSession("admin", ...)`
|
||||||
|
- [x] 2.3 新建 `src/web/pages/auth/UserLoginPage.tsx`，覆盖 `pages/UserLogin.tsx` 当前逻辑，写入会话调 `commitSession("user", ...)`；底部预留一个 `<OAuthSlot />` 占位组件，本次不挂任何 provider 按钮，留给后续 change
|
||||||
|
- [x] 2.4 把 `src/web/App.tsx` 中 `Login` 组件替换为 `<AdminLoginPage>`，删除组件内部的 useState/handler/JSX；保留 onLogin 回调签名
|
||||||
|
- [x] 2.5 把 `src/web/pages/UserLogin.tsx` 改为薄壳：`export { default } from './auth/UserLoginPage'` 或 1 行包裹组件
|
||||||
|
- [x] 2.6 迁移 `App.login-surface.test.tsx` 中 DOM 断言到 `src/web/pages/auth/LoginShell.test.tsx` 与 `AdminLoginPage.test.tsx` / `UserLoginPage.test.tsx`；保留 `App.login-surface.test.tsx` 仅断言挂载入口存在
|
||||||
|
- [x] 2.7 跑 `npm run typecheck:web && npm run test`，确认无回归
|
||||||
|
|
||||||
|
## 3. 清理与验收
|
||||||
|
|
||||||
|
- [x] 3.1 删除 `App.tsx` 中已无引用的 helper / 文案常量；用 `rg "import.*from.*authSession"` 巡检无遗漏调用 `persistAuthSession` / `persistUserSession` 的页面
|
||||||
|
- [x] 3.2 更新 `README.md` 与 `AGENTS.md` 中"登录"段落，加入新目录结构说明（`pages/auth/`）；OAuth 登录入口标注"由后续 change 接入"
|
||||||
|
- [x] 3.3 跑全量校验：`npm run typecheck` / `npm run test` / `cargo test --lib` / `cargo clippy --all-targets`
|
||||||
|
- typecheck:web、typecheck:web:test 全过；typecheck:server 因 `tsconfig.server.json` 不存在报 TS5058，stash 校验确认 pre-existing
|
||||||
|
- npm run test：本 change 把 baseline 的 11 files / 17 tests 失败减到 9 files / 13 tests（修好 `App.mobile-layout` 与 `App.sidebar-mobile`），剩余 9 个失败文件均与 authSession 无关、stash 校验确认 pre-existing
|
||||||
|
- cargo test --lib：78 passed / 3 failed；3 个失败 (`proxy::token_router::router::tests::resolves_exact_model_route` / `…resolves_wildcard_route_and_skips_excluded_channel` / `routes::proxy::tests::chat_route_uses_token_router_and_conductor`) 同样在 stash 后基线复现，pre-existing
|
||||||
|
- cargo clippy --all-targets：本 change 区块零警告（`comfyui_agent.rs` 测试 mutex 已切到 `tokio::sync::Mutex`）；仅余一处 unrelated `clippy::needless_borrow`（`token_router/router.rs:289` 一带）pre-existing
|
||||||
|
- [ ] 3.4 手测三条登录路径（管理员 token、管理员账密、用户账密）；OAuth 占位区域确保渲染但不挂行为
|
||||||
|
- [x] 3.5 在 PR 描述中列出所有 spec scenarios 的对应测试位置，方便 reviewer 校对（mapping 见下文「附：scenario → test 映射」）
|
||||||
|
|
||||||
|
## 附：scenario → test 映射
|
||||||
|
|
||||||
|
### specs/frontend-login/spec.md
|
||||||
|
|
||||||
|
- 登录界面拆分为独立模块
|
||||||
|
- 管理员入口挂载新组件 → `src/web/pages/auth/AdminLoginPage.test.tsx`（断言挂载来源 + 路由 `/`）
|
||||||
|
- 用户入口保留旧路由 → `src/web/pages/auth/UserLoginPage.test.tsx` + `src/web/pages/UserLogin.tsx` 静态包含断言（薄壳）
|
||||||
|
- App.tsx 不再持有登录表单状态 → `src/web/App.login-surface.test.tsx`（保留为入口断言）+ rg 巡检（CI 时由 grep 守卫脚本兜底）
|
||||||
|
- 统一登录错误解析
|
||||||
|
- 401 / 网络异常 / 400+reason → `src/web/loginError.test.ts` 三组用例
|
||||||
|
- 会话写入收敛为单一 API
|
||||||
|
- 管理员登录成功调用 commitSession → `src/web/pages/auth/AdminLoginPage.test.tsx`（mock authSession，断言 commitSession('admin', token) 被调用）
|
||||||
|
- 用户登录成功调用 commitSession → `src/web/pages/auth/UserLoginPage.test.tsx`
|
||||||
|
- 旧 API 仍可工作 → `src/web/authSession.test.ts`（薄壳兼容用例）
|
||||||
|
- 登录页预留 OAuth 占位槽
|
||||||
|
- 用户登录页渲染占位槽 → `src/web/pages/auth/UserLoginPage.test.tsx`（断言 `data-testid="oauth-slot"` 存在 + 无 fetch）
|
||||||
|
- 管理员登录页不含占位槽 → `src/web/pages/auth/AdminLoginPage.test.tsx`（反向断言）
|
||||||
|
- 占位槽不触发任何行为 → `src/web/pages/auth/OAuthSlot.test.tsx`（点击无 fetch / 无 commit / 无 navigate）
|
||||||
@@ -0,0 +1,20 @@
|
|||||||
|
schema: spec-driven
|
||||||
|
|
||||||
|
# Project context (optional)
|
||||||
|
# This is shown to AI when creating artifacts.
|
||||||
|
# Add your tech stack, conventions, style guides, domain knowledge, etc.
|
||||||
|
# Example:
|
||||||
|
# context: |
|
||||||
|
# Tech stack: TypeScript, React, Node.js
|
||||||
|
# We use conventional commits
|
||||||
|
# Domain: e-commerce platform
|
||||||
|
|
||||||
|
# Per-artifact rules (optional)
|
||||||
|
# Add custom rules for specific artifacts.
|
||||||
|
# Example:
|
||||||
|
# rules:
|
||||||
|
# proposal:
|
||||||
|
# - Keep proposals under 500 words
|
||||||
|
# - Always include a "Non-goals" section
|
||||||
|
# tasks:
|
||||||
|
# - Break tasks into chunks of max 2 hours
|
||||||
@@ -0,0 +1,92 @@
|
|||||||
|
# route-candidate-api Specification
|
||||||
|
|
||||||
|
## Purpose
|
||||||
|
TBD - created by archiving change add-route-rebuild-pipeline. Update Purpose after archive.
|
||||||
|
## Requirements
|
||||||
|
### Requirement: `GET /api/models/token-candidates` is registered
|
||||||
|
|
||||||
|
The admin router SHALL register `GET /api/models/token-candidates`
|
||||||
|
behind the standard admin auth layer. The handler SHALL return a JSON
|
||||||
|
object with exactly four top-level keys: `models`,
|
||||||
|
`modelsWithoutToken`, `modelsMissingTokenGroups`, `endpointTypesByModel`.
|
||||||
|
|
||||||
|
#### Scenario: Endpoint returns 200 with the four-key shape
|
||||||
|
- **WHEN** an authenticated admin sends `GET /api/models/token-candidates`
|
||||||
|
- **THEN** the response is HTTP 200 with a JSON body containing keys
|
||||||
|
`models`, `modelsWithoutToken`, `modelsMissingTokenGroups`, and
|
||||||
|
`endpointTypesByModel` — no extra top-level keys, no missing keys
|
||||||
|
|
||||||
|
#### Scenario: Unauthenticated request is rejected
|
||||||
|
- **WHEN** a request without a valid admin token reaches the endpoint
|
||||||
|
- **THEN** the admin auth layer returns 401 before the handler runs
|
||||||
|
|
||||||
|
### Requirement: `models` lists every (model, token) candidate
|
||||||
|
|
||||||
|
The `models` field SHALL be an object whose keys are model names and
|
||||||
|
whose values are arrays of candidate descriptors. A `(model, token)`
|
||||||
|
combination SHALL appear in the `models` array for that model exactly
|
||||||
|
when ALL of the following hold:
|
||||||
|
- a row in `model_availability` reports `available = 1` for the model
|
||||||
|
on the token's account
|
||||||
|
- the token's `enabled` is truthy (NULL treated as enabled)
|
||||||
|
- the token's `value_status` is `'ready'` (NULL treated as ready)
|
||||||
|
- the account's `status = 'active'`
|
||||||
|
- the site's `status` is `'active'` (NULL treated as active)
|
||||||
|
|
||||||
|
Each candidate descriptor SHALL contain: `accountId`, `tokenId`,
|
||||||
|
`tokenName`, `isDefault` (boolean), `username`, `siteId`, `siteName`.
|
||||||
|
Within a single model's array, descriptors SHALL be deduplicated by
|
||||||
|
`tokenId`.
|
||||||
|
|
||||||
|
#### Scenario: Active token with available model appears
|
||||||
|
- **WHEN** account 5 (active, site active) has token 11 (enabled,
|
||||||
|
ready) and `model_availability(account=5, model='gpt-4')` is
|
||||||
|
available
|
||||||
|
- **THEN** `models['gpt-4']` includes a descriptor with
|
||||||
|
`accountId = 5, tokenId = 11`
|
||||||
|
|
||||||
|
#### Scenario: Disabled token is excluded
|
||||||
|
- **WHEN** account 5 has token 11 with `enabled = 0`
|
||||||
|
- **THEN** no descriptor with `tokenId = 11` appears in `models` for
|
||||||
|
any model
|
||||||
|
|
||||||
|
#### Scenario: Duplicate tokens in source data deduped
|
||||||
|
- **WHEN** the underlying join produces two rows for the same
|
||||||
|
`(model, tokenId)` pair (e.g., joins fan out)
|
||||||
|
- **THEN** the descriptor appears exactly once in
|
||||||
|
`models[model_name]`
|
||||||
|
|
||||||
|
### Requirement: `modelsWithoutToken` surfaces accounts missing tokens
|
||||||
|
|
||||||
|
The `modelsWithoutToken` field SHALL be an object keyed by model name.
|
||||||
|
For each `(account, model)` combination where the model is available
|
||||||
|
on the active account but no qualifying token (per the `models`
|
||||||
|
requirement) exists for that account, the handler SHALL include a
|
||||||
|
descriptor `{ accountId, username, siteId, siteName }` under that
|
||||||
|
model name. Descriptors SHALL be deduplicated by `accountId` within
|
||||||
|
each model's array.
|
||||||
|
|
||||||
|
#### Scenario: Account with available model but no enabled tokens
|
||||||
|
- **WHEN** account 8 reports `gpt-4` available but has zero rows in
|
||||||
|
`account_tokens` (or all rows are disabled)
|
||||||
|
- **THEN** `modelsWithoutToken['gpt-4']` includes a descriptor with
|
||||||
|
`accountId = 8`, and that account does NOT appear in `models['gpt-4']`
|
||||||
|
|
||||||
|
#### Scenario: Account with token does not appear here
|
||||||
|
- **WHEN** account 5 reports `gpt-4` available and has a usable token
|
||||||
|
- **THEN** account 5 does NOT appear in `modelsWithoutToken['gpt-4']`
|
||||||
|
|
||||||
|
### Requirement: `modelsMissingTokenGroups` and `endpointTypesByModel` are empty
|
||||||
|
|
||||||
|
The handler SHALL return `modelsMissingTokenGroups` and
|
||||||
|
`endpointTypesByModel` as empty JSON objects (`{}`). The frontend
|
||||||
|
contract permits empty objects and these fields SHALL NOT be
|
||||||
|
populated until a marketplace-cache or pricing-catalog source is
|
||||||
|
added to the backend.
|
||||||
|
|
||||||
|
#### Scenario: Both fields are empty objects
|
||||||
|
- **WHEN** the handler returns successfully under any state of the
|
||||||
|
database
|
||||||
|
- **THEN** `response.modelsMissingTokenGroups === {}` and
|
||||||
|
`response.endpointTypesByModel === {}`
|
||||||
|
|
||||||
@@ -0,0 +1,157 @@
|
|||||||
|
# route-rebuild-pipeline Specification
|
||||||
|
|
||||||
|
## Purpose
|
||||||
|
TBD - created by archiving change add-route-rebuild-pipeline. Update Purpose after archive.
|
||||||
|
## Requirements
|
||||||
|
### Requirement: Route rebuild reads from `model_availability`
|
||||||
|
|
||||||
|
The rebuild pipeline SHALL derive its working set from the
|
||||||
|
`model_availability` table joined to `accounts`, filtering out:
|
||||||
|
- accounts whose `status` is not `'active'`
|
||||||
|
- rows whose `available` is not `1` (treating `NULL` as not available)
|
||||||
|
- `(account, model)` combinations covered by `site_disabled_models`
|
||||||
|
for the account's `site_id`
|
||||||
|
|
||||||
|
The pipeline SHALL NOT make outbound HTTP requests to upstream platforms
|
||||||
|
during a rebuild. Refreshing `model_availability` is the responsibility
|
||||||
|
of the existing `POST /api/accounts/{id}/models/refresh` endpoint and
|
||||||
|
is out of scope for this requirement.
|
||||||
|
|
||||||
|
#### Scenario: Inactive account is excluded
|
||||||
|
- **WHEN** an account with `status = 'suspended'` has rows in
|
||||||
|
`model_availability`
|
||||||
|
- **THEN** no `token_routes` row is created on its behalf and no
|
||||||
|
`route_channels` row references its `account_id`
|
||||||
|
|
||||||
|
#### Scenario: Site-disabled model is excluded
|
||||||
|
- **WHEN** `site_disabled_models` contains `(site_id = 7, model_name = 'gpt-x')`
|
||||||
|
and an active account on site 7 has `model_availability` for `gpt-x`
|
||||||
|
- **THEN** the rebuild does not include that account in the channel set
|
||||||
|
for the `gpt-x` route
|
||||||
|
|
||||||
|
### Requirement: Rebuild aggregates one route per distinct model name
|
||||||
|
|
||||||
|
The rebuild SHALL ensure exactly one row exists in `token_routes` for
|
||||||
|
every distinct `model_name` in the filtered working set, keyed by
|
||||||
|
`model_pattern = model_name`. The rebuild SHALL reuse existing rows
|
||||||
|
when their `model_pattern` already matches; new rows SHALL be inserted
|
||||||
|
with `routing_strategy = 'weighted'`, `route_mode = 'pattern'`, and
|
||||||
|
`enabled = 1`.
|
||||||
|
|
||||||
|
The rebuild MUST NOT overwrite `display_name`, `display_icon`, or any
|
||||||
|
other fields on existing `token_routes` rows. Operators rely on these
|
||||||
|
fields being editable from the frontend without being clobbered on
|
||||||
|
re-runs.
|
||||||
|
|
||||||
|
#### Scenario: New model creates a new route
|
||||||
|
- **WHEN** model `claude-4-opus` appears in the working set and no
|
||||||
|
`token_routes` row has `model_pattern = 'claude-4-opus'`
|
||||||
|
- **THEN** the rebuild inserts a `token_routes` row with
|
||||||
|
`model_pattern = 'claude-4-opus'`, `routing_strategy = 'weighted'`,
|
||||||
|
`route_mode = 'pattern'`, `enabled = 1`
|
||||||
|
|
||||||
|
#### Scenario: Existing route is reused without overwrite
|
||||||
|
- **WHEN** a `token_routes` row exists with
|
||||||
|
`model_pattern = 'gpt-4'` and `display_name = 'GPT-4 Turbo (Custom)'`
|
||||||
|
- **THEN** the rebuild leaves `display_name` unchanged and reuses the
|
||||||
|
same `id` for channel reconciliation
|
||||||
|
|
||||||
|
### Requirement: Channel reconciliation preserves manual overrides
|
||||||
|
|
||||||
|
For each route, the rebuild SHALL reconcile `route_channels` rows
|
||||||
|
keyed by `(account_id, source_model)` against the desired set derived
|
||||||
|
from `model_availability`. Rows with `manual_override = 1` SHALL be
|
||||||
|
treated as immutable: not inserted, not updated, not deleted. Rows
|
||||||
|
with `manual_override = 0` (or `NULL`) SHALL be:
|
||||||
|
- inserted when missing from the existing set
|
||||||
|
- hard-deleted when no longer in the desired set
|
||||||
|
|
||||||
|
The desired set SHALL exclude any `(account, model)` combination
|
||||||
|
already covered by a `manual_override = 1` row, so the rebuild never
|
||||||
|
inserts an auto channel that shadows a manual one.
|
||||||
|
|
||||||
|
#### Scenario: Manual channel survives rebuild
|
||||||
|
- **WHEN** a `route_channels` row has `manual_override = 1` and its
|
||||||
|
`(account_id, source_model)` is no longer in the desired set
|
||||||
|
- **THEN** the row is left untouched and counted in
|
||||||
|
`RebuildSummary.skipped_manual_channels`
|
||||||
|
|
||||||
|
#### Scenario: Stale auto channel is deleted
|
||||||
|
- **WHEN** a `route_channels` row has `manual_override = 0`,
|
||||||
|
`account_id = 12`, `source_model = 'old-model'`, and no longer
|
||||||
|
appears in the desired set
|
||||||
|
- **THEN** the row is deleted and counted in
|
||||||
|
`RebuildSummary.deleted_channels`
|
||||||
|
|
||||||
|
#### Scenario: New auto channel is inserted with default token
|
||||||
|
- **WHEN** account 12 newly reports `gpt-5` available and has at least
|
||||||
|
one enabled `account_tokens` row
|
||||||
|
- **THEN** a new `route_channels` row is inserted with `route_id`
|
||||||
|
pointing at the `gpt-5` route, `account_id = 12`, the
|
||||||
|
default-flagged enabled token's id (or the lowest id when none is
|
||||||
|
default), `priority = 0`, `weight = 10`, `enabled = 1`,
|
||||||
|
`manual_override = 0`
|
||||||
|
|
||||||
|
#### Scenario: Account without tokens still gets a channel
|
||||||
|
- **WHEN** account 13 reports `gpt-5` available but has no enabled
|
||||||
|
`account_tokens` rows
|
||||||
|
- **THEN** a `route_channels` row is inserted with `token_id = NULL`,
|
||||||
|
relying on the proxy layer's account-level token fallback
|
||||||
|
|
||||||
|
### Requirement: Rebuild runs in a single transaction
|
||||||
|
|
||||||
|
The rebuild SHALL execute every read and write inside one database
|
||||||
|
transaction, committed at the end. A partial failure MUST NOT leave
|
||||||
|
`token_routes` and `route_channels` mutually inconsistent.
|
||||||
|
|
||||||
|
#### Scenario: Failure mid-rebuild rolls back
|
||||||
|
- **WHEN** an `INSERT` into `route_channels` fails partway through
|
||||||
|
the rebuild
|
||||||
|
- **THEN** the transaction aborts and no `token_routes` insert from
|
||||||
|
the same rebuild persists
|
||||||
|
|
||||||
|
### Requirement: Rebuild summary matches frontend contract
|
||||||
|
|
||||||
|
The rebuild SHALL return a `RebuildSummary` containing the fields
|
||||||
|
`createdRoutes`, `reusedRoutes`, `createdChannels`, `deletedChannels`,
|
||||||
|
`skippedManualChannels`, `scannedModels`, `scannedAccounts` (camelCase
|
||||||
|
when serialized). The `POST /api/routes/rebuild` handler SHALL embed
|
||||||
|
this object under the JSON key `rebuild`, alongside `success: true`.
|
||||||
|
|
||||||
|
`TokenRoutes.tsx:411-412` reads `rebuild.createdRoutes` and
|
||||||
|
`rebuild.createdChannels`; both MUST be present and SHALL count only
|
||||||
|
rows actually inserted by the current rebuild (not pre-existing rows
|
||||||
|
that were reused).
|
||||||
|
|
||||||
|
#### Scenario: Rebuild response shape
|
||||||
|
- **WHEN** the handler completes successfully
|
||||||
|
- **THEN** the response body is
|
||||||
|
`{ "success": true, "rebuild": { "createdRoutes": <n1>,
|
||||||
|
"reusedRoutes": <n2>, "createdChannels": <n3>,
|
||||||
|
"deletedChannels": <n4>, "skippedManualChannels": <n5>,
|
||||||
|
"scannedModels": <n6>, "scannedAccounts": <n7> } }`
|
||||||
|
|
||||||
|
#### Scenario: Reused route does not increment createdRoutes
|
||||||
|
- **WHEN** every model in the working set has a pre-existing
|
||||||
|
`token_routes` row
|
||||||
|
- **THEN** `createdRoutes = 0` and `reusedRoutes` equals the number of
|
||||||
|
distinct models scanned
|
||||||
|
|
||||||
|
### Requirement: Auto-rebuild after account creation
|
||||||
|
|
||||||
|
The system SHALL fire one route rebuild in the background after
|
||||||
|
`POST /api/accounts` returns success for at least one created
|
||||||
|
account. The HTTP response SHALL return without waiting for the
|
||||||
|
rebuild to finish. Errors from the background rebuild SHALL be logged
|
||||||
|
at `warn` level and SHALL NOT affect the response status or body.
|
||||||
|
|
||||||
|
#### Scenario: Single-account create triggers rebuild
|
||||||
|
- **WHEN** `POST /api/accounts` creates exactly one account
|
||||||
|
- **THEN** `spawn_route_rebuild` is invoked once with the cloned
|
||||||
|
`AppState` and the response returns immediately
|
||||||
|
|
||||||
|
#### Scenario: Batch-create with all failures does not trigger
|
||||||
|
- **WHEN** `POST /api/accounts` returns with `created.len() == 0` and
|
||||||
|
`failed.len() > 0`
|
||||||
|
- **THEN** `spawn_route_rebuild` is not invoked
|
||||||
|
|
||||||
@@ -24,6 +24,12 @@ pub enum AppError {
|
|||||||
#[error("Rate limited")]
|
#[error("Rate limited")]
|
||||||
RateLimited,
|
RateLimited,
|
||||||
|
|
||||||
|
#[error("Service unavailable: {0}")]
|
||||||
|
ServiceUnavailable(String),
|
||||||
|
|
||||||
|
#[error("Gateway timeout: {0}")]
|
||||||
|
GatewayTimeout(String),
|
||||||
|
|
||||||
#[error("Internal error: {0}")]
|
#[error("Internal error: {0}")]
|
||||||
Internal(#[from] anyhow::Error),
|
Internal(#[from] anyhow::Error),
|
||||||
|
|
||||||
@@ -40,6 +46,8 @@ impl IntoResponse for AppError {
|
|||||||
AppError::BadRequest(msg) => (StatusCode::BAD_REQUEST, msg.clone()),
|
AppError::BadRequest(msg) => (StatusCode::BAD_REQUEST, msg.clone()),
|
||||||
AppError::UpstreamError(msg) => (StatusCode::BAD_GATEWAY, msg.clone()),
|
AppError::UpstreamError(msg) => (StatusCode::BAD_GATEWAY, msg.clone()),
|
||||||
AppError::RateLimited => (StatusCode::TOO_MANY_REQUESTS, "rate_limited".into()),
|
AppError::RateLimited => (StatusCode::TOO_MANY_REQUESTS, "rate_limited".into()),
|
||||||
|
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::Internal(e) => (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()),
|
||||||
AppError::Database(e) => (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()),
|
AppError::Database(e) => (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()),
|
||||||
};
|
};
|
||||||
|
|||||||
+12
-1
@@ -16,6 +16,7 @@ pub struct AppState {
|
|||||||
pub config: config::AppConfig,
|
pub config: config::AppConfig,
|
||||||
pub db: Arc<SqlitePool>,
|
pub db: Arc<SqlitePool>,
|
||||||
pub comfyui: Arc<services::comfyui_process::ComfyUIProcess>,
|
pub comfyui: Arc<services::comfyui_process::ComfyUIProcess>,
|
||||||
|
pub comfyui_pool: Arc<services::comfyui_pool::ComfyUIPool>,
|
||||||
pub builtin_keys: Arc<services::builtin_keys::BuiltinKeys>,
|
pub builtin_keys: Arc<services::builtin_keys::BuiltinKeys>,
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -34,15 +35,25 @@ impl AppState {
|
|||||||
let comfyui_agent_secret = builtin_keys.comfyui_agent.1.clone();
|
let comfyui_agent_secret = builtin_keys.comfyui_agent.1.clone();
|
||||||
let comfyui = services::comfyui_process::ComfyUIProcess::new(
|
let comfyui = services::comfyui_process::ComfyUIProcess::new(
|
||||||
config.comfyui.clone(),
|
config.comfyui.clone(),
|
||||||
boosapi_url,
|
boosapi_url.clone(),
|
||||||
config.auth.proxy_token.clone(),
|
config.auth.proxy_token.clone(),
|
||||||
|
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,
|
comfyui_agent_secret,
|
||||||
);
|
);
|
||||||
|
let comfyui_pool = services::comfyui_pool::ComfyUIPool::new(docker, pool_cfg);
|
||||||
|
|
||||||
Self {
|
Self {
|
||||||
config,
|
config,
|
||||||
db: Arc::new(db),
|
db: Arc::new(db),
|
||||||
comfyui,
|
comfyui,
|
||||||
|
comfyui_pool,
|
||||||
builtin_keys: Arc::new(builtin_keys),
|
builtin_keys: Arc::new(builtin_keys),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
+21
@@ -105,6 +105,20 @@ async fn main() -> anyhow::Result<()> {
|
|||||||
}
|
}
|
||||||
let comfyui_process = state.comfyui.clone();
|
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
|
// Build router
|
||||||
let app = boosapi::routes::build_router(state);
|
let app = boosapi::routes::build_router(state);
|
||||||
|
|
||||||
@@ -123,6 +137,13 @@ async fn main() -> anyhow::Result<()> {
|
|||||||
tracing::warn!(%err, "ComfyUI embedded process did not stop cleanly");
|
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(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -209,3 +209,80 @@ pub async fn execute_image_generation(
|
|||||||
|
|
||||||
Ok(json)
|
Ok(json)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// One field forwarded to the upstream `/images/edits` endpoint.
|
||||||
|
///
|
||||||
|
/// `bytes` is None → text field; Some → file field. We don't try to
|
||||||
|
/// detect the difference at the HTTP layer because OpenAI's spec already
|
||||||
|
/// pins which fields are files (`image`, `image[]`, `mask`) — the route
|
||||||
|
/// handler classifies them based on the field's content-type/filename.
|
||||||
|
pub struct MultipartField {
|
||||||
|
pub name: String,
|
||||||
|
pub bytes: Option<Vec<u8>>,
|
||||||
|
pub filename: Option<String>,
|
||||||
|
pub content_type: Option<String>,
|
||||||
|
/// Text value for non-file fields. Ignored when `bytes` is Some.
|
||||||
|
pub text: Option<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Execute an image edits request (`POST /v1/images/edits`).
|
||||||
|
///
|
||||||
|
/// Unlike /images/generations this is `multipart/form-data` (the request
|
||||||
|
/// carries one or more image binaries plus the prompt). We rebuild a
|
||||||
|
/// fresh upstream multipart form from the parsed fields rather than
|
||||||
|
/// streaming the raw body — that lets us swap out the `model` field with
|
||||||
|
/// the route-resolved upstream model and skip any client-side fields
|
||||||
|
/// that don't belong on the wire.
|
||||||
|
pub async fn execute_image_edit(
|
||||||
|
_state: AppState,
|
||||||
|
route: UpstreamRoute,
|
||||||
|
fields: Vec<MultipartField>,
|
||||||
|
) -> AppResult<Value> {
|
||||||
|
use reqwest::multipart::{Form, Part};
|
||||||
|
|
||||||
|
let client = http_client();
|
||||||
|
let upstream_url = format!("{}/images/edits", route.base_url.trim_end_matches('/'));
|
||||||
|
|
||||||
|
let mut form = Form::new();
|
||||||
|
for field in fields {
|
||||||
|
if let Some(bytes) = field.bytes {
|
||||||
|
let mut part = Part::bytes(bytes);
|
||||||
|
if let Some(filename) = field.filename {
|
||||||
|
part = part.file_name(filename);
|
||||||
|
}
|
||||||
|
if let Some(ct) = field.content_type {
|
||||||
|
part = part
|
||||||
|
.mime_str(&ct)
|
||||||
|
.map_err(|e| AppError::BadRequest(format!("invalid content-type: {e}")))?;
|
||||||
|
}
|
||||||
|
form = form.part(field.name, part);
|
||||||
|
} else {
|
||||||
|
form = form.text(field.name, field.text.unwrap_or_default());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let resp = client
|
||||||
|
.post(&upstream_url)
|
||||||
|
.header("Authorization", format!("Bearer {}", route.api_key))
|
||||||
|
.multipart(form)
|
||||||
|
.send()
|
||||||
|
.await
|
||||||
|
.map_err(|e| AppError::UpstreamError(format!("Upstream request failed: {e}")))?;
|
||||||
|
|
||||||
|
let status = resp.status();
|
||||||
|
let text = resp
|
||||||
|
.text()
|
||||||
|
.await
|
||||||
|
.map_err(|e| AppError::UpstreamError(format!("Failed to read upstream response: {e}")))?;
|
||||||
|
|
||||||
|
if !status.is_success() {
|
||||||
|
return Err(AppError::UpstreamError(format!(
|
||||||
|
"Upstream returned {status}: {text}"
|
||||||
|
)));
|
||||||
|
}
|
||||||
|
|
||||||
|
let json: Value = serde_json::from_str(&text)
|
||||||
|
.map_err(|e| AppError::UpstreamError(format!("Invalid JSON from upstream: {e}")))?;
|
||||||
|
|
||||||
|
Ok(json)
|
||||||
|
}
|
||||||
|
|||||||
@@ -376,6 +376,22 @@ mod tests {
|
|||||||
.execute(&pool)
|
.execute(&pool)
|
||||||
.await
|
.await
|
||||||
.unwrap();
|
.unwrap();
|
||||||
|
// select_channel_for_route's COALESCE subquery joins through
|
||||||
|
// site_api_endpoints. The table must exist or the SELECT errors
|
||||||
|
// out — we leave it empty and the COALESCE falls back to sites.url.
|
||||||
|
sqlx::query(
|
||||||
|
"CREATE TABLE site_api_endpoints (
|
||||||
|
id INTEGER PRIMARY KEY,
|
||||||
|
site_id INTEGER NOT NULL,
|
||||||
|
url TEXT NOT NULL,
|
||||||
|
enabled INTEGER DEFAULT 1,
|
||||||
|
sort_order INTEGER DEFAULT 0,
|
||||||
|
cooldown_until TEXT
|
||||||
|
)",
|
||||||
|
)
|
||||||
|
.execute(&pool)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
sqlx::query("INSERT INTO sites (id, name, url, api_key) VALUES (1, 'OpenAI', 'https://api.example/v1', 'site-key')")
|
sqlx::query("INSERT INTO sites (id, name, url, api_key) VALUES (1, 'OpenAI', 'https://api.example/v1', 'site-key')")
|
||||||
.execute(&pool)
|
.execute(&pool)
|
||||||
|
|||||||
+162
-11
@@ -147,16 +147,8 @@ pub async fn list_accounts(
|
|||||||
let enriched: Vec<Value> = accounts
|
let enriched: Vec<Value> = accounts
|
||||||
.iter()
|
.iter()
|
||||||
.map(|account| {
|
.map(|account| {
|
||||||
let mut row = serde_json::to_value(account).unwrap_or(json!({}));
|
let site = site_by_id.get(&account.site_id).copied();
|
||||||
if let Some(site) = site_by_id.get(&account.site_id) {
|
account_response_value(account, site)
|
||||||
if let Value::Object(map) = &mut row {
|
|
||||||
map.insert(
|
|
||||||
"site".to_string(),
|
|
||||||
serde_json::to_value(site).unwrap_or(Value::Null),
|
|
||||||
);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
row
|
|
||||||
})
|
})
|
||||||
.collect();
|
.collect();
|
||||||
|
|
||||||
@@ -168,6 +160,49 @@ pub async fn list_accounts(
|
|||||||
})))
|
})))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn account_response_value(account: &Account, site: Option<&Site>) -> Value {
|
||||||
|
let credential_mode = resolve_account_credential_mode(account);
|
||||||
|
let mut row = serde_json::to_value(account).unwrap_or(json!({}));
|
||||||
|
if let Value::Object(map) = &mut row {
|
||||||
|
map.insert("credentialMode".to_string(), json!(credential_mode));
|
||||||
|
map.insert(
|
||||||
|
"capabilities".to_string(),
|
||||||
|
json!({
|
||||||
|
"canCheckin": credential_mode == "session",
|
||||||
|
"canRefreshBalance": credential_mode == "session",
|
||||||
|
"proxyOnly": credential_mode == "apikey",
|
||||||
|
}),
|
||||||
|
);
|
||||||
|
if let Some(site) = site {
|
||||||
|
map.insert(
|
||||||
|
"site".to_string(),
|
||||||
|
serde_json::to_value(site).unwrap_or(Value::Null),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
row
|
||||||
|
}
|
||||||
|
|
||||||
|
fn resolve_account_credential_mode(account: &Account) -> &'static str {
|
||||||
|
let Some(extra_config) = account.extra_config.as_deref() else {
|
||||||
|
return "apikey";
|
||||||
|
};
|
||||||
|
let Ok(value) = serde_json::from_str::<Value>(extra_config) else {
|
||||||
|
return "apikey";
|
||||||
|
};
|
||||||
|
match value
|
||||||
|
.get("credentialMode")
|
||||||
|
.and_then(Value::as_str)
|
||||||
|
.unwrap_or("")
|
||||||
|
.trim()
|
||||||
|
.to_ascii_lowercase()
|
||||||
|
.as_str()
|
||||||
|
{
|
||||||
|
"session" => "session",
|
||||||
|
_ => "apikey",
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
/// POST /api/accounts
|
/// POST /api/accounts
|
||||||
pub async fn create_account(
|
pub async fn create_account(
|
||||||
State(state): State<AppState>,
|
State(state): State<AppState>,
|
||||||
@@ -279,12 +314,17 @@ pub async fn create_account(
|
|||||||
|
|
||||||
spawn_route_rebuild(state.clone());
|
spawn_route_rebuild(state.clone());
|
||||||
|
|
||||||
|
let site: Option<Site> = sqlx::query_as::<_, Site>("SELECT * FROM sites WHERE id = ?")
|
||||||
|
.bind(account.site_id)
|
||||||
|
.fetch_optional(&*state.db)
|
||||||
|
.await?;
|
||||||
|
|
||||||
Ok((
|
Ok((
|
||||||
StatusCode::CREATED,
|
StatusCode::CREATED,
|
||||||
Json(json!({
|
Json(json!({
|
||||||
"success": true,
|
"success": true,
|
||||||
"id": account.id,
|
"id": account.id,
|
||||||
"data": account,
|
"data": account_response_value(&account, site.as_ref()),
|
||||||
})),
|
})),
|
||||||
))
|
))
|
||||||
}
|
}
|
||||||
@@ -524,6 +564,117 @@ pub async fn delete_account(
|
|||||||
Ok(Json(json!({ "data": { "deleted": true, "id": id } })))
|
Ok(Json(json!({ "data": { "deleted": true, "id": id } })))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Deserialize)]
|
||||||
|
pub struct BatchAccountsRequest {
|
||||||
|
pub ids: Vec<i64>,
|
||||||
|
pub action: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// POST /api/accounts/batch
|
||||||
|
///
|
||||||
|
/// Per-id loop, no transaction: a delete on row 5 failing must not roll back
|
||||||
|
/// the successful delete on row 3 — the frontend's `runBatchAccountAction`
|
||||||
|
/// reads `successIds` / `failedItems` and re-selects only the failed ones so
|
||||||
|
/// the operator can retry the bad rows in isolation.
|
||||||
|
///
|
||||||
|
/// Supported actions:
|
||||||
|
/// * `delete` — `DELETE FROM accounts WHERE id = ?`
|
||||||
|
/// * `enable` — `status = 'active'`
|
||||||
|
/// * `disable` — `status = 'disabled'`
|
||||||
|
/// * `refreshBalance` — not implemented; every id goes to `failedItems` so
|
||||||
|
/// the UI shows "成功 0, 失败 N" instead of silently
|
||||||
|
/// claiming success without changing the balance row.
|
||||||
|
///
|
||||||
|
/// `spawn_route_rebuild` is fired once at the end if any row was actually
|
||||||
|
/// touched (status flip or delete) — `rebuild_routes` filters on
|
||||||
|
/// `status = 'active'`, so any change to that set invalidates the route_channels
|
||||||
|
/// view and needs a rebuild. Doing it once after the loop (rather than per-id)
|
||||||
|
/// avoids N concurrent rebuilds racing each other.
|
||||||
|
pub async fn batch_accounts(
|
||||||
|
State(state): State<AppState>,
|
||||||
|
Json(payload): Json<BatchAccountsRequest>,
|
||||||
|
) -> AppResult<Json<Value>> {
|
||||||
|
info!(
|
||||||
|
"Batch account action: action={} count={}",
|
||||||
|
payload.action,
|
||||||
|
payload.ids.len()
|
||||||
|
);
|
||||||
|
|
||||||
|
if payload.ids.is_empty() {
|
||||||
|
return Ok(Json(json!({
|
||||||
|
"data": { "successIds": [], "failedItems": [] }
|
||||||
|
})));
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut success_ids: Vec<i64> = Vec::with_capacity(payload.ids.len());
|
||||||
|
let mut failed_items: Vec<Value> = Vec::new();
|
||||||
|
let mut routing_dirty = false;
|
||||||
|
|
||||||
|
for id in &payload.ids {
|
||||||
|
let result: Result<(), String> = match payload.action.as_str() {
|
||||||
|
"delete" => {
|
||||||
|
match sqlx::query("DELETE FROM accounts WHERE id = ?")
|
||||||
|
.bind(id)
|
||||||
|
.execute(&*state.db)
|
||||||
|
.await
|
||||||
|
{
|
||||||
|
Ok(r) if r.rows_affected() == 0 => Err("账号不存在".into()),
|
||||||
|
Ok(_) => {
|
||||||
|
routing_dirty = true;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
Err(e) => Err(format!("DB error: {}", e)),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
"enable" | "disable" => {
|
||||||
|
let new_status = if payload.action == "enable" {
|
||||||
|
"active"
|
||||||
|
} else {
|
||||||
|
"disabled"
|
||||||
|
};
|
||||||
|
match sqlx::query(
|
||||||
|
"UPDATE accounts SET status = ?, updated_at = datetime('now') WHERE id = ?",
|
||||||
|
)
|
||||||
|
.bind(new_status)
|
||||||
|
.bind(id)
|
||||||
|
.execute(&*state.db)
|
||||||
|
.await
|
||||||
|
{
|
||||||
|
Ok(r) if r.rows_affected() == 0 => Err("账号不存在".into()),
|
||||||
|
Ok(_) => {
|
||||||
|
routing_dirty = true;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
Err(e) => Err(format!("DB error: {}", e)),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
"refreshBalance" => {
|
||||||
|
// Per-site balance probing isn't implemented yet — see proposal
|
||||||
|
// notes. Mark every id failed so the UI surfaces "成功 0" rather
|
||||||
|
// than pretending the balance was refreshed.
|
||||||
|
Err("单站点余额探测尚未实现".into())
|
||||||
|
}
|
||||||
|
other => Err(format!("不支持的批量操作: {}", other)),
|
||||||
|
};
|
||||||
|
|
||||||
|
match result {
|
||||||
|
Ok(()) => success_ids.push(*id),
|
||||||
|
Err(msg) => failed_items.push(json!({ "id": id, "error": msg })),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if routing_dirty {
|
||||||
|
spawn_route_rebuild(state.clone());
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(Json(json!({
|
||||||
|
"data": {
|
||||||
|
"successIds": success_ids,
|
||||||
|
"failedItems": failed_items,
|
||||||
|
}
|
||||||
|
})))
|
||||||
|
}
|
||||||
|
|
||||||
async fn resolve_site_api_base_url(state: &AppState, site: &Site) -> AppResult<String> {
|
async fn resolve_site_api_base_url(state: &AppState, site: &Site) -> AppResult<String> {
|
||||||
let endpoint: Option<String> = sqlx::query_scalar(
|
let endpoint: Option<String> = sqlx::query_scalar(
|
||||||
"SELECT url
|
"SELECT url
|
||||||
|
|||||||
@@ -365,8 +365,12 @@ async fn clear_route_cooldown(
|
|||||||
Ok(Json(json!({ "success": true })))
|
Ok(Json(json!({ "success": true })))
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn rebuild_routes() -> Json<Value> {
|
async fn rebuild_routes(State(state): State<AppState>) -> AppResult<Json<Value>> {
|
||||||
Json(json!({ "success": true, "message": "routes already available" }))
|
let summary = crate::services::route_rebuild::rebuild_routes(&state).await?;
|
||||||
|
Ok(Json(json!({
|
||||||
|
"success": true,
|
||||||
|
"rebuild": summary,
|
||||||
|
})))
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn batch_routes(
|
async fn batch_routes(
|
||||||
|
|||||||
@@ -13,6 +13,38 @@ use serde_json::{json, Value};
|
|||||||
use tracing::info;
|
use tracing::info;
|
||||||
use uuid::Uuid;
|
use uuid::Uuid;
|
||||||
|
|
||||||
|
/// Sentinel name for the built-in ComfyUI Agent downstream key.
|
||||||
|
///
|
||||||
|
/// The system row is uniquely identified by `name = SYSTEM_KEY_NAME AND
|
||||||
|
/// user_id IS NULL` — see `services::builtin_keys::ensure_comfyui_agent_key`.
|
||||||
|
const SYSTEM_KEY_NAME: &str = "comfyui-agent";
|
||||||
|
|
||||||
|
fn is_system_key(name: &str, user_id: Option<i64>) -> bool {
|
||||||
|
name == SYSTEM_KEY_NAME && user_id.is_none()
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Mask a key down to `••••••••<last4>`. Operates on chars so non-ASCII
|
||||||
|
/// keys (unlikely but possible) do not panic on byte slicing.
|
||||||
|
fn mask_key(key: &str) -> String {
|
||||||
|
let chars: Vec<char> = key.chars().collect();
|
||||||
|
let last4: String = chars[chars.len().saturating_sub(4)..].iter().collect();
|
||||||
|
format!("••••••••{}", last4)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Project a `DownstreamApiKey` row to its admin-list JSON shape, masking the
|
||||||
|
/// secret and tagging `system: true` for the built-in `comfyui-agent` row.
|
||||||
|
fn key_to_response(key: DownstreamApiKey) -> Value {
|
||||||
|
let system = is_system_key(&key.name, key.user_id);
|
||||||
|
let mut value = serde_json::to_value(&key).unwrap_or_else(|_| json!({}));
|
||||||
|
if system {
|
||||||
|
if let Value::Object(ref mut map) = value {
|
||||||
|
map.insert("system".to_string(), Value::Bool(true));
|
||||||
|
map.insert("key".to_string(), Value::String(mask_key(&key.key)));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
value
|
||||||
|
}
|
||||||
|
|
||||||
#[derive(Debug, Deserialize)]
|
#[derive(Debug, Deserialize)]
|
||||||
pub struct ListKeysQuery {
|
pub struct ListKeysQuery {
|
||||||
pub q: Option<String>,
|
pub q: Option<String>,
|
||||||
@@ -89,8 +121,9 @@ pub async fn list_downstream_keys(
|
|||||||
.fetch_all(&*state.db)
|
.fetch_all(&*state.db)
|
||||||
.await?;
|
.await?;
|
||||||
|
|
||||||
|
let projected: Vec<Value> = keys.into_iter().map(key_to_response).collect();
|
||||||
Ok(Json(
|
Ok(Json(
|
||||||
json!({ "success": true, "items": keys, "data": keys }),
|
json!({ "success": true, "items": projected, "data": projected }),
|
||||||
))
|
))
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -170,6 +203,16 @@ pub async fn update_downstream_key(
|
|||||||
AppError::NotFound(format!("Downstream key with id {} not found", id))
|
AppError::NotFound(format!("Downstream key with id {} not found", id))
|
||||||
})?;
|
})?;
|
||||||
|
|
||||||
|
// Guard: refuse any request that tries to disable the built-in
|
||||||
|
// `comfyui-agent` system key (`name='comfyui-agent' AND user_id IS NULL`).
|
||||||
|
// The agent's outbound calls in Local mode bear this key — disabling it
|
||||||
|
// would silently break image / TTS / chat for every user.
|
||||||
|
if is_system_key(&existing.name, existing.user_id) && payload.enabled == Some(false) {
|
||||||
|
return Err(AppError::Forbidden(
|
||||||
|
"system key cannot be disabled".into(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
|
||||||
let supported_models: Option<String> = if let Some(ref models) = payload.supported_models {
|
let supported_models: Option<String> = if let Some(ref models) = payload.supported_models {
|
||||||
validate_supported_models(models)?;
|
validate_supported_models(models)?;
|
||||||
Some(
|
Some(
|
||||||
@@ -214,6 +257,24 @@ pub async fn delete_downstream_key(
|
|||||||
) -> AppResult<Json<Value>> {
|
) -> AppResult<Json<Value>> {
|
||||||
info!("Deleting downstream API key: {}", id);
|
info!("Deleting downstream API key: {}", id);
|
||||||
|
|
||||||
|
// Pre-fetch so we can refuse the built-in `comfyui-agent` row before
|
||||||
|
// running the DELETE. Returning 403 instead of silently filtering keeps
|
||||||
|
// the failure observable to operators who attempt the action.
|
||||||
|
let existing: DownstreamApiKey =
|
||||||
|
sqlx::query_as::<_, DownstreamApiKey>("SELECT * FROM downstream_api_keys WHERE id = ?")
|
||||||
|
.bind(id)
|
||||||
|
.fetch_optional(&*state.db)
|
||||||
|
.await?
|
||||||
|
.ok_or_else(|| {
|
||||||
|
AppError::NotFound(format!("Downstream key with id {} not found", id))
|
||||||
|
})?;
|
||||||
|
|
||||||
|
if is_system_key(&existing.name, existing.user_id) {
|
||||||
|
return Err(AppError::Forbidden(
|
||||||
|
"system key cannot be deleted".into(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
|
||||||
let affected = sqlx::query("DELETE FROM downstream_api_keys WHERE id = ?")
|
let affected = sqlx::query("DELETE FROM downstream_api_keys WHERE id = ?")
|
||||||
.bind(id)
|
.bind(id)
|
||||||
.execute(&*state.db)
|
.execute(&*state.db)
|
||||||
@@ -266,3 +327,213 @@ pub async fn reset_downstream_key(
|
|||||||
|
|
||||||
Ok(Json(json!({ "data": api_key })))
|
Ok(Json(json!({ "data": api_key })))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
use crate::config::{AppConfig, AuthConfig, DbConfig, ProxyConfig, ServerConfig};
|
||||||
|
use crate::services::builtin_keys::BuiltinKeys;
|
||||||
|
use sqlx::sqlite::SqlitePoolOptions;
|
||||||
|
|
||||||
|
async fn test_state() -> AppState {
|
||||||
|
let pool = SqlitePoolOptions::new()
|
||||||
|
.max_connections(1)
|
||||||
|
.connect("sqlite::memory:")
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
sqlx::query(
|
||||||
|
"CREATE TABLE downstream_api_keys (
|
||||||
|
id INTEGER PRIMARY KEY,
|
||||||
|
name TEXT NOT NULL,
|
||||||
|
key TEXT NOT NULL,
|
||||||
|
description TEXT,
|
||||||
|
group_name TEXT,
|
||||||
|
tags TEXT,
|
||||||
|
enabled INTEGER DEFAULT 1,
|
||||||
|
expires_at TEXT,
|
||||||
|
max_cost REAL,
|
||||||
|
used_cost REAL DEFAULT 0,
|
||||||
|
max_requests INTEGER,
|
||||||
|
used_requests INTEGER 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 INTEGER,
|
||||||
|
created_at TEXT DEFAULT (datetime('now')),
|
||||||
|
updated_at TEXT DEFAULT (datetime('now'))
|
||||||
|
)",
|
||||||
|
)
|
||||||
|
.execute(&pool)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
// Row 1: built-in system key (name='comfyui-agent', user_id IS NULL)
|
||||||
|
sqlx::query(
|
||||||
|
"INSERT INTO downstream_api_keys (id, name, key, enabled, user_id) VALUES (1, 'comfyui-agent', 'sk-comfyui-abcdef0123456789', 1, NULL)"
|
||||||
|
)
|
||||||
|
.execute(&pool)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
// Row 2: regular admin-scoped key
|
||||||
|
sqlx::query(
|
||||||
|
"INSERT INTO downstream_api_keys (id, name, key, enabled, user_id) VALUES (2, 'team-prod', 'sk-prod-xyz123', 1, NULL)"
|
||||||
|
)
|
||||||
|
.execute(&pool)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
AppState::new(
|
||||||
|
AppConfig {
|
||||||
|
server: ServerConfig {
|
||||||
|
host: "127.0.0.1".into(),
|
||||||
|
port: 0,
|
||||||
|
log_level: "info".into(),
|
||||||
|
cors_origins: vec![],
|
||||||
|
},
|
||||||
|
db: DbConfig {
|
||||||
|
url: "sqlite::memory:".into(),
|
||||||
|
max_connections: 1,
|
||||||
|
ssl: false,
|
||||||
|
runtime_db_url: None,
|
||||||
|
runtime_db_dialect: None,
|
||||||
|
runtime_db_ssl: None,
|
||||||
|
},
|
||||||
|
auth: AuthConfig {
|
||||||
|
jwt_secret: "secret".into(),
|
||||||
|
jwt_expiry_hours: 1,
|
||||||
|
auth_token: "auth".into(),
|
||||||
|
proxy_token: "proxy".into(),
|
||||||
|
},
|
||||||
|
proxy: ProxyConfig {
|
||||||
|
upstream_timeout_secs: 30,
|
||||||
|
first_byte_timeout_secs: 30,
|
||||||
|
max_retries: 3,
|
||||||
|
proxy_log_retention_days: 7,
|
||||||
|
},
|
||||||
|
comfyui: Default::default(),
|
||||||
|
},
|
||||||
|
pool,
|
||||||
|
BuiltinKeys {
|
||||||
|
comfyui_agent: (1, "sk-comfyui-abcdef0123456789".into()),
|
||||||
|
},
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 4.2 — Listing: the system row carries `system: true` and a masked
|
||||||
|
/// `key`, while regular rows pass through unchanged.
|
||||||
|
#[tokio::test]
|
||||||
|
async fn list_masks_system_key_and_passes_regular_keys_through() {
|
||||||
|
let state = test_state().await;
|
||||||
|
let resp = list_downstream_keys(
|
||||||
|
State(state),
|
||||||
|
Query(ListKeysQuery {
|
||||||
|
q: None,
|
||||||
|
user_id: None,
|
||||||
|
}),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.expect("list should succeed");
|
||||||
|
let items = resp.0.get("items").and_then(Value::as_array).unwrap();
|
||||||
|
assert_eq!(items.len(), 2);
|
||||||
|
|
||||||
|
let system = items
|
||||||
|
.iter()
|
||||||
|
.find(|item| item["id"] == 1)
|
||||||
|
.expect("system row");
|
||||||
|
assert_eq!(system["system"], Value::Bool(true));
|
||||||
|
assert_eq!(system["key"], "••••••••6789");
|
||||||
|
assert_eq!(system["name"], "comfyui-agent");
|
||||||
|
|
||||||
|
let regular = items
|
||||||
|
.iter()
|
||||||
|
.find(|item| item["id"] == 2)
|
||||||
|
.expect("regular row");
|
||||||
|
assert!(
|
||||||
|
regular.get("system").is_none(),
|
||||||
|
"regular keys must not carry the system flag"
|
||||||
|
);
|
||||||
|
assert_eq!(regular["key"], "sk-prod-xyz123");
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 4.3 — DELETE on the built-in row returns 403 with the expected message.
|
||||||
|
#[tokio::test]
|
||||||
|
async fn delete_rejects_system_key_with_forbidden() {
|
||||||
|
let state = test_state().await;
|
||||||
|
let result = delete_downstream_key(State(state), Path(1)).await;
|
||||||
|
match result {
|
||||||
|
Err(AppError::Forbidden(msg)) => assert_eq!(msg, "system key cannot be deleted"),
|
||||||
|
other => panic!("expected Forbidden, got {:?}", other),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 4.5 — DELETE on a regular row succeeds (control case).
|
||||||
|
#[tokio::test]
|
||||||
|
async fn delete_succeeds_for_regular_key() {
|
||||||
|
let state = test_state().await;
|
||||||
|
let resp = delete_downstream_key(State(state.clone()), Path(2))
|
||||||
|
.await
|
||||||
|
.expect("delete should succeed for regular key");
|
||||||
|
assert_eq!(resp.0["data"]["deleted"], Value::Bool(true));
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 4.4 — Setting `enabled: false` on the system row returns 403; the row
|
||||||
|
/// remains enabled in the DB.
|
||||||
|
#[tokio::test]
|
||||||
|
async fn update_rejects_disabling_system_key() {
|
||||||
|
let state = test_state().await;
|
||||||
|
let payload = UpdateDownstreamKeyRequest {
|
||||||
|
name: None,
|
||||||
|
description: None,
|
||||||
|
group_name: None,
|
||||||
|
tags: None,
|
||||||
|
enabled: Some(false),
|
||||||
|
expires_at: None,
|
||||||
|
max_cost: None,
|
||||||
|
max_requests: None,
|
||||||
|
supported_models: None,
|
||||||
|
allowed_route_ids: None,
|
||||||
|
user_id: None,
|
||||||
|
};
|
||||||
|
let result = update_downstream_key(State(state.clone()), Path(1), Json(payload)).await;
|
||||||
|
match result {
|
||||||
|
Err(AppError::Forbidden(msg)) => assert_eq!(msg, "system key cannot be disabled"),
|
||||||
|
other => panic!("expected Forbidden, got {:?}", other),
|
||||||
|
}
|
||||||
|
|
||||||
|
// Confirm the row was not touched.
|
||||||
|
let still_enabled: Option<bool> =
|
||||||
|
sqlx::query_scalar("SELECT enabled FROM downstream_api_keys WHERE id = 1")
|
||||||
|
.fetch_one(&*state.db)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
assert_eq!(still_enabled, Some(true));
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 4.4 — Other field updates on the system row (e.g. `description`) are
|
||||||
|
/// allowed; only `enabled = false` is gated.
|
||||||
|
#[tokio::test]
|
||||||
|
async fn update_allows_non_disable_changes_on_system_key() {
|
||||||
|
let state = test_state().await;
|
||||||
|
let payload = UpdateDownstreamKeyRequest {
|
||||||
|
name: None,
|
||||||
|
description: Some("internal: do not delete".into()),
|
||||||
|
group_name: None,
|
||||||
|
tags: None,
|
||||||
|
enabled: None,
|
||||||
|
expires_at: None,
|
||||||
|
max_cost: None,
|
||||||
|
max_requests: None,
|
||||||
|
supported_models: None,
|
||||||
|
allowed_route_ids: None,
|
||||||
|
user_id: None,
|
||||||
|
};
|
||||||
|
let resp = update_downstream_key(State(state), Path(1), Json(payload))
|
||||||
|
.await
|
||||||
|
.expect("non-disable update should pass");
|
||||||
|
assert_eq!(resp.0["data"]["description"], "internal: do not delete");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
+172
-2
@@ -18,6 +18,8 @@ mod account_tokens;
|
|||||||
mod accounts;
|
mod accounts;
|
||||||
mod compat;
|
mod compat;
|
||||||
mod downstream_keys;
|
mod downstream_keys;
|
||||||
|
mod model_candidates;
|
||||||
|
mod models_marketplace;
|
||||||
mod sites;
|
mod sites;
|
||||||
mod tokens;
|
mod tokens;
|
||||||
mod users;
|
mod users;
|
||||||
@@ -66,6 +68,11 @@ pub fn routes(state: AppState) -> Router<AppState> {
|
|||||||
get(sites::get_site)
|
get(sites::get_site)
|
||||||
.put(sites::update_site)
|
.put(sites::update_site)
|
||||||
.delete(sites::delete_site),
|
.delete(sites::delete_site),
|
||||||
|
)
|
||||||
|
.route(
|
||||||
|
"/api/sites/{id}/disabled-models",
|
||||||
|
get(sites::list_site_disabled_models)
|
||||||
|
.put(sites::update_site_disabled_models),
|
||||||
);
|
);
|
||||||
|
|
||||||
let accounts_routes = Router::new()
|
let accounts_routes = Router::new()
|
||||||
@@ -73,6 +80,7 @@ pub fn routes(state: AppState) -> Router<AppState> {
|
|||||||
"/api/accounts",
|
"/api/accounts",
|
||||||
get(accounts::list_accounts).post(accounts::create_account),
|
get(accounts::list_accounts).post(accounts::create_account),
|
||||||
)
|
)
|
||||||
|
.route("/api/accounts/batch", post(accounts::batch_accounts))
|
||||||
.route("/api/accounts/verify-token", post(accounts::verify_token))
|
.route("/api/accounts/verify-token", post(accounts::verify_token))
|
||||||
.route(
|
.route(
|
||||||
"/api/accounts/{id}",
|
"/api/accounts/{id}",
|
||||||
@@ -87,6 +95,10 @@ pub fn routes(state: AppState) -> Router<AppState> {
|
|||||||
.route(
|
.route(
|
||||||
"/api/accounts/{id}/models/manual",
|
"/api/accounts/{id}/models/manual",
|
||||||
post(accounts::add_account_manual_models),
|
post(accounts::add_account_manual_models),
|
||||||
|
)
|
||||||
|
.route(
|
||||||
|
"/api/accounts/{id}/models/refresh",
|
||||||
|
post(accounts::refresh_account_models),
|
||||||
);
|
);
|
||||||
|
|
||||||
let tokens_routes = Router::new()
|
let tokens_routes = Router::new()
|
||||||
@@ -137,6 +149,16 @@ pub fn routes(state: AppState) -> Router<AppState> {
|
|||||||
let settings_routes =
|
let settings_routes =
|
||||||
Router::new().route("/api/settings", get(get_settings).put(update_setting));
|
Router::new().route("/api/settings", get(get_settings).put(update_setting));
|
||||||
|
|
||||||
|
let model_candidates_routes = Router::new().route(
|
||||||
|
"/api/models/token-candidates",
|
||||||
|
get(model_candidates::token_candidates),
|
||||||
|
);
|
||||||
|
|
||||||
|
let models_marketplace_routes = Router::new().route(
|
||||||
|
"/api/models/marketplace",
|
||||||
|
get(models_marketplace::get_marketplace),
|
||||||
|
);
|
||||||
|
|
||||||
let users_routes = Router::new()
|
let users_routes = Router::new()
|
||||||
.route(
|
.route(
|
||||||
"/api/admin/users",
|
"/api/admin/users",
|
||||||
@@ -170,6 +192,8 @@ pub fn routes(state: AppState) -> Router<AppState> {
|
|||||||
.merge(downstream_routes)
|
.merge(downstream_routes)
|
||||||
.merge(settings_routes)
|
.merge(settings_routes)
|
||||||
.merge(users_routes)
|
.merge(users_routes)
|
||||||
|
.merge(model_candidates_routes)
|
||||||
|
.merge(models_marketplace_routes)
|
||||||
.merge(account_tokens::routes())
|
.merge(account_tokens::routes())
|
||||||
.merge(compat::routes())
|
.merge(compat::routes())
|
||||||
.route_layer(middleware::from_fn_with_state(state, require_admin_auth))
|
.route_layer(middleware::from_fn_with_state(state, require_admin_auth))
|
||||||
@@ -177,7 +201,20 @@ pub fn routes(state: AppState) -> Router<AppState> {
|
|||||||
|
|
||||||
// ── Settings handlers ────────────────────────────────────────────────
|
// ── Settings handlers ────────────────────────────────────────────────
|
||||||
|
|
||||||
|
/// Settings rows whose keys begin with this prefix are managed by
|
||||||
|
/// `services::builtin_keys` and contain server-issued secrets (e.g. the
|
||||||
|
/// `comfyui-agent` downstream key). They must never appear in admin
|
||||||
|
/// list/get/edit responses.
|
||||||
|
const HIDDEN_SETTINGS_PREFIX: &str = "builtin_key.";
|
||||||
|
|
||||||
|
fn is_hidden_setting_key(key: &str) -> bool {
|
||||||
|
key.starts_with(HIDDEN_SETTINGS_PREFIX)
|
||||||
|
}
|
||||||
|
|
||||||
/// GET /api/settings — return all settings as a key-value map.
|
/// GET /api/settings — return all settings as a key-value map.
|
||||||
|
///
|
||||||
|
/// Filters out any keys with the `builtin_key.` prefix so the admin UI
|
||||||
|
/// never sees secrets bootstrapped by `services::builtin_keys`.
|
||||||
async fn get_settings(State(state): State<AppState>) -> AppResult<Json<Value>> {
|
async fn get_settings(State(state): State<AppState>) -> AppResult<Json<Value>> {
|
||||||
info!("Getting all settings");
|
info!("Getting all settings");
|
||||||
|
|
||||||
@@ -186,14 +223,21 @@ async fn get_settings(State(state): State<AppState>) -> AppResult<Json<Value>> {
|
|||||||
.fetch_all(&*state.db)
|
.fetch_all(&*state.db)
|
||||||
.await?;
|
.await?;
|
||||||
|
|
||||||
let map: std::collections::BTreeMap<String, Option<String>> =
|
let map: std::collections::BTreeMap<String, Option<String>> = rows
|
||||||
rows.into_iter().map(|s| (s.key, s.value)).collect();
|
.into_iter()
|
||||||
|
.filter(|s| !is_hidden_setting_key(&s.key))
|
||||||
|
.map(|s| (s.key, s.value))
|
||||||
|
.collect();
|
||||||
|
|
||||||
Ok(Json(json!({ "data": map })))
|
Ok(Json(json!({ "data": map })))
|
||||||
}
|
}
|
||||||
|
|
||||||
/// PUT /api/settings — upsert a setting by key.
|
/// PUT /api/settings — upsert a setting by key.
|
||||||
/// Expects JSON body: { "key": "...", "value": "..." }
|
/// Expects JSON body: { "key": "...", "value": "..." }
|
||||||
|
///
|
||||||
|
/// Refuses any key with the `builtin_key.` prefix: those rows are owned by
|
||||||
|
/// the bootstrap path and a stray PUT would clobber the secret used by the
|
||||||
|
/// ComfyUI Agent's outbound calls.
|
||||||
async fn update_setting(
|
async fn update_setting(
|
||||||
State(state): State<AppState>,
|
State(state): State<AppState>,
|
||||||
Json(payload): Json<serde_json::Value>,
|
Json(payload): Json<serde_json::Value>,
|
||||||
@@ -203,6 +247,12 @@ async fn update_setting(
|
|||||||
.and_then(|k| k.as_str())
|
.and_then(|k| k.as_str())
|
||||||
.ok_or_else(|| AppError::BadRequest("key is required".into()))?;
|
.ok_or_else(|| AppError::BadRequest("key is required".into()))?;
|
||||||
|
|
||||||
|
if is_hidden_setting_key(key) {
|
||||||
|
return Err(AppError::Forbidden(
|
||||||
|
"system setting cannot be modified".into(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
|
||||||
let value = payload.get("value").and_then(|v| {
|
let value = payload.get("value").and_then(|v| {
|
||||||
if v.is_null() {
|
if v.is_null() {
|
||||||
None
|
None
|
||||||
@@ -226,3 +276,123 @@ async fn update_setting(
|
|||||||
}
|
}
|
||||||
})))
|
})))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
use crate::config::{AppConfig, AuthConfig, DbConfig, ProxyConfig, ServerConfig};
|
||||||
|
use crate::services::builtin_keys::BuiltinKeys;
|
||||||
|
use sqlx::sqlite::SqlitePoolOptions;
|
||||||
|
|
||||||
|
async fn settings_state() -> AppState {
|
||||||
|
let pool = SqlitePoolOptions::new()
|
||||||
|
.max_connections(1)
|
||||||
|
.connect("sqlite::memory:")
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
sqlx::query(
|
||||||
|
"CREATE TABLE settings (key TEXT PRIMARY KEY, value TEXT)",
|
||||||
|
)
|
||||||
|
.execute(&pool)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
sqlx::query("INSERT INTO settings (key, value) VALUES ('site.title', 'BoosAPI')")
|
||||||
|
.execute(&pool)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
sqlx::query(
|
||||||
|
"INSERT INTO settings (key, value) VALUES ('builtin_key.comfyui_agent.secret', 'sk-comfyui-secret')",
|
||||||
|
)
|
||||||
|
.execute(&pool)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
AppState::new(
|
||||||
|
AppConfig {
|
||||||
|
server: ServerConfig {
|
||||||
|
host: "127.0.0.1".into(),
|
||||||
|
port: 0,
|
||||||
|
log_level: "info".into(),
|
||||||
|
cors_origins: vec![],
|
||||||
|
},
|
||||||
|
db: DbConfig {
|
||||||
|
url: "sqlite::memory:".into(),
|
||||||
|
max_connections: 1,
|
||||||
|
ssl: false,
|
||||||
|
runtime_db_url: None,
|
||||||
|
runtime_db_dialect: None,
|
||||||
|
runtime_db_ssl: None,
|
||||||
|
},
|
||||||
|
auth: AuthConfig {
|
||||||
|
jwt_secret: "secret".into(),
|
||||||
|
jwt_expiry_hours: 1,
|
||||||
|
auth_token: "auth".into(),
|
||||||
|
proxy_token: "proxy".into(),
|
||||||
|
},
|
||||||
|
proxy: ProxyConfig {
|
||||||
|
upstream_timeout_secs: 30,
|
||||||
|
first_byte_timeout_secs: 30,
|
||||||
|
max_retries: 3,
|
||||||
|
proxy_log_retention_days: 7,
|
||||||
|
},
|
||||||
|
comfyui: Default::default(),
|
||||||
|
},
|
||||||
|
pool,
|
||||||
|
BuiltinKeys {
|
||||||
|
comfyui_agent: (1, "sk-comfyui-secret".into()),
|
||||||
|
},
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 5.2 — `GET /api/settings` returns regular settings but never any
|
||||||
|
/// `builtin_key.*` rows.
|
||||||
|
#[tokio::test]
|
||||||
|
async fn list_settings_filters_builtin_key_prefix() {
|
||||||
|
let state = settings_state().await;
|
||||||
|
let resp = get_settings(State(state)).await.expect("list ok");
|
||||||
|
let data = resp.0.get("data").and_then(Value::as_object).unwrap();
|
||||||
|
assert!(data.contains_key("site.title"));
|
||||||
|
assert!(
|
||||||
|
!data.contains_key("builtin_key.comfyui_agent.secret"),
|
||||||
|
"builtin_key.* rows must not appear in /api/settings response"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 5.3 — `PUT /api/settings` rejects writes targeting `builtin_key.*`
|
||||||
|
/// keys with 403; the row is left untouched.
|
||||||
|
#[tokio::test]
|
||||||
|
async fn update_settings_rejects_builtin_key_writes() {
|
||||||
|
let state = settings_state().await;
|
||||||
|
let payload = json!({
|
||||||
|
"key": "builtin_key.comfyui_agent.secret",
|
||||||
|
"value": "sk-attacker"
|
||||||
|
});
|
||||||
|
let result = update_setting(State(state.clone()), Json(payload)).await;
|
||||||
|
match result {
|
||||||
|
Err(AppError::Forbidden(msg)) => {
|
||||||
|
assert_eq!(msg, "system setting cannot be modified");
|
||||||
|
}
|
||||||
|
other => panic!("expected Forbidden, got {:?}", other),
|
||||||
|
}
|
||||||
|
|
||||||
|
let stored: Option<String> = sqlx::query_scalar(
|
||||||
|
"SELECT value FROM settings WHERE key = 'builtin_key.comfyui_agent.secret'",
|
||||||
|
)
|
||||||
|
.fetch_one(&*state.db)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
assert_eq!(stored.as_deref(), Some("sk-comfyui-secret"));
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 5.3 — Regular settings still upsert as before.
|
||||||
|
#[tokio::test]
|
||||||
|
async fn update_settings_allows_regular_keys() {
|
||||||
|
let state = settings_state().await;
|
||||||
|
let payload = json!({ "key": "site.title", "value": "Renamed" });
|
||||||
|
let resp = update_setting(State(state.clone()), Json(payload))
|
||||||
|
.await
|
||||||
|
.expect("update ok");
|
||||||
|
assert_eq!(resp.0["data"]["value"], "Renamed");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -0,0 +1,176 @@
|
|||||||
|
//! `GET /api/models/token-candidates` — feeds the frontend "candidate models"
|
||||||
|
//! panel and the new-route dialog dropdown (`TokenRoutes.tsx:252`).
|
||||||
|
//!
|
||||||
|
//! Response shape (must match the legacy Node handler):
|
||||||
|
//!
|
||||||
|
//! ```json
|
||||||
|
//! {
|
||||||
|
//! "models": {
|
||||||
|
//! "<modelName>": [
|
||||||
|
//! { "accountId", "tokenId", "tokenName", "isDefault",
|
||||||
|
//! "username", "siteId", "siteName" }, ...
|
||||||
|
//! ]
|
||||||
|
//! },
|
||||||
|
//! "modelsWithoutToken": {
|
||||||
|
//! "<modelName>": [ { "accountId", "username", "siteId", "siteName" }, ... ]
|
||||||
|
//! },
|
||||||
|
//! "modelsMissingTokenGroups": {},
|
||||||
|
//! "endpointTypesByModel": {}
|
||||||
|
//! }
|
||||||
|
//! ```
|
||||||
|
//!
|
||||||
|
//! `modelsMissingTokenGroups` and `endpointTypesByModel` are intentionally empty:
|
||||||
|
//! this codebase doesn't have the pricing-catalog metadata or
|
||||||
|
//! marketplace-cache machinery the Node version relies on. Returning empty
|
||||||
|
//! objects keeps the contract intact (frontend tolerates them) without
|
||||||
|
//! pulling in unrelated infrastructure.
|
||||||
|
|
||||||
|
use crate::{error::AppResult, AppState};
|
||||||
|
use axum::{extract::State, Json};
|
||||||
|
use serde_json::{json, Map, Value};
|
||||||
|
use std::collections::{HashMap, HashSet};
|
||||||
|
|
||||||
|
/// Row from `model_availability ⋈ account_tokens ⋈ accounts ⋈ sites` — every
|
||||||
|
/// `(model, token)` combo that's currently usable.
|
||||||
|
#[derive(Debug, sqlx::FromRow)]
|
||||||
|
struct CoveredRow {
|
||||||
|
model_name: String,
|
||||||
|
account_id: i64,
|
||||||
|
token_id: i64,
|
||||||
|
token_name: String,
|
||||||
|
is_default: Option<bool>,
|
||||||
|
username: Option<String>,
|
||||||
|
site_id: i64,
|
||||||
|
site_name: Option<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Row from `model_availability ⋈ accounts ⋈ sites` — every `(model, account)`
|
||||||
|
/// combo regardless of token coverage. We diff this against `CoveredRow` to
|
||||||
|
/// surface accounts that have the model but no usable token.
|
||||||
|
#[derive(Debug, sqlx::FromRow)]
|
||||||
|
struct AccountModelRow {
|
||||||
|
model_name: String,
|
||||||
|
account_id: i64,
|
||||||
|
username: Option<String>,
|
||||||
|
site_id: i64,
|
||||||
|
site_name: Option<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn token_candidates(State(state): State<AppState>) -> AppResult<Json<Value>> {
|
||||||
|
// 1. Models that already have at least one usable token. The joins drop:
|
||||||
|
// inactive accounts/sites, disabled tokens, and tokens whose value
|
||||||
|
// isn't ready (e.g. mid-rotation).
|
||||||
|
let covered: Vec<CoveredRow> = sqlx::query_as(
|
||||||
|
"SELECT ma.model_name AS model_name,
|
||||||
|
a.id AS account_id,
|
||||||
|
at.id AS token_id,
|
||||||
|
at.name AS token_name,
|
||||||
|
at.is_default AS is_default,
|
||||||
|
a.username AS username,
|
||||||
|
s.id AS site_id,
|
||||||
|
s.name AS site_name
|
||||||
|
FROM model_availability ma
|
||||||
|
JOIN account_tokens at ON at.account_id = ma.account_id
|
||||||
|
JOIN accounts a ON a.id = ma.account_id
|
||||||
|
JOIN sites s ON s.id = a.site_id
|
||||||
|
WHERE COALESCE(ma.available, 0) = 1
|
||||||
|
AND COALESCE(at.enabled, 1) = 1
|
||||||
|
AND COALESCE(at.value_status, 'ready') = 'ready'
|
||||||
|
AND a.status = 'active'
|
||||||
|
AND COALESCE(s.status, 'active') = 'active'",
|
||||||
|
)
|
||||||
|
.fetch_all(&*state.db)
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
// 2. All available `(account, model)` pairs — used to compute which
|
||||||
|
// accounts have the model but no usable token (`modelsWithoutToken`).
|
||||||
|
let account_models: Vec<AccountModelRow> = sqlx::query_as(
|
||||||
|
"SELECT ma.model_name AS model_name,
|
||||||
|
a.id AS account_id,
|
||||||
|
a.username AS username,
|
||||||
|
s.id AS site_id,
|
||||||
|
s.name AS site_name
|
||||||
|
FROM model_availability ma
|
||||||
|
JOIN accounts a ON a.id = ma.account_id
|
||||||
|
JOIN sites s ON s.id = a.site_id
|
||||||
|
WHERE COALESCE(ma.available, 0) = 1
|
||||||
|
AND a.status = 'active'
|
||||||
|
AND COALESCE(s.status, 'active') = 'active'",
|
||||||
|
)
|
||||||
|
.fetch_all(&*state.db)
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
// Build `models` map. Dedupe candidates by tokenId per model so a token
|
||||||
|
// marked available against the same model twice doesn't render twice.
|
||||||
|
let mut models: HashMap<String, Vec<Value>> = HashMap::new();
|
||||||
|
let mut seen_token_per_model: HashMap<String, HashSet<i64>> = HashMap::new();
|
||||||
|
let mut covered_account_model: HashSet<(i64, String)> = HashSet::new();
|
||||||
|
|
||||||
|
for row in covered {
|
||||||
|
let model_name = row.model_name.trim().to_string();
|
||||||
|
if model_name.is_empty() {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
covered_account_model.insert((row.account_id, model_name.to_lowercase()));
|
||||||
|
|
||||||
|
let seen = seen_token_per_model.entry(model_name.clone()).or_default();
|
||||||
|
if !seen.insert(row.token_id) {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
models.entry(model_name).or_default().push(json!({
|
||||||
|
"accountId": row.account_id,
|
||||||
|
"tokenId": row.token_id,
|
||||||
|
"tokenName": row.token_name,
|
||||||
|
"isDefault": row.is_default.unwrap_or(false),
|
||||||
|
"username": row.username,
|
||||||
|
"siteId": row.site_id,
|
||||||
|
"siteName": row.site_name,
|
||||||
|
}));
|
||||||
|
}
|
||||||
|
|
||||||
|
// Build `modelsWithoutToken`: every (account, model) in `account_models`
|
||||||
|
// that didn't show up in `covered`. Dedupe by accountId per model.
|
||||||
|
let mut models_without_token: HashMap<String, Vec<Value>> = HashMap::new();
|
||||||
|
let mut seen_account_per_model: HashMap<String, HashSet<i64>> = HashMap::new();
|
||||||
|
for row in account_models {
|
||||||
|
let model_name = row.model_name.trim().to_string();
|
||||||
|
if model_name.is_empty() {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
let key = (row.account_id, model_name.to_lowercase());
|
||||||
|
if covered_account_model.contains(&key) {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
let seen = seen_account_per_model
|
||||||
|
.entry(model_name.clone())
|
||||||
|
.or_default();
|
||||||
|
if !seen.insert(row.account_id) {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
models_without_token.entry(model_name).or_default().push(json!({
|
||||||
|
"accountId": row.account_id,
|
||||||
|
"username": row.username,
|
||||||
|
"siteId": row.site_id,
|
||||||
|
"siteName": row.site_name,
|
||||||
|
}));
|
||||||
|
}
|
||||||
|
|
||||||
|
// Convert HashMap -> serde_json::Map so the response keys serialize
|
||||||
|
// deterministically as a JSON object (HashMap also works but Map is the
|
||||||
|
// canonical type for object-shaped payloads).
|
||||||
|
let models_obj: Map<String, Value> = models
|
||||||
|
.into_iter()
|
||||||
|
.map(|(k, v)| (k, Value::Array(v)))
|
||||||
|
.collect();
|
||||||
|
let without_obj: Map<String, Value> = models_without_token
|
||||||
|
.into_iter()
|
||||||
|
.map(|(k, v)| (k, Value::Array(v)))
|
||||||
|
.collect();
|
||||||
|
|
||||||
|
Ok(Json(json!({
|
||||||
|
"models": models_obj,
|
||||||
|
"modelsWithoutToken": without_obj,
|
||||||
|
"modelsMissingTokenGroups": {},
|
||||||
|
"endpointTypesByModel": {},
|
||||||
|
})))
|
||||||
|
}
|
||||||
@@ -0,0 +1,379 @@
|
|||||||
|
//! `GET /api/models/marketplace` — model marketplace listing for the admin
|
||||||
|
//! `/models` page (frontend: `Models.tsx`).
|
||||||
|
//!
|
||||||
|
//! Stage 1: aggregates `model_availability` × `accounts` × `account_tokens`
|
||||||
|
//! × `sites` into the shape `ModelsMarketplaceResponse` expects. Per-model
|
||||||
|
//! latency / success rate / description / tags / supportedEndpointTypes /
|
||||||
|
//! pricingSources are returned as null/empty placeholders — the frontend
|
||||||
|
//! falls back gracefully and stage 2 will hydrate them from `proxy_logs`
|
||||||
|
//! and pricing tables.
|
||||||
|
//!
|
||||||
|
//! `?refresh=1` and `?includePricing=1` are accepted but currently no-ops:
|
||||||
|
//! we always return the in-memory aggregation and an empty `meta`. Stage 2
|
||||||
|
//! will wire them to an async refresh job and the pricing hydration path.
|
||||||
|
|
||||||
|
use crate::{error::AppResult, AppState};
|
||||||
|
use axum::{
|
||||||
|
extract::{Query, State},
|
||||||
|
Json,
|
||||||
|
};
|
||||||
|
use serde::Deserialize;
|
||||||
|
use serde_json::{json, Value};
|
||||||
|
use sqlx::FromRow;
|
||||||
|
use std::collections::{BTreeMap, HashMap};
|
||||||
|
|
||||||
|
#[derive(Debug, Deserialize)]
|
||||||
|
#[allow(dead_code)] // stage 2 will read these; keep the fields to lock the API shape.
|
||||||
|
pub struct MarketplaceQuery {
|
||||||
|
pub refresh: Option<String>,
|
||||||
|
#[serde(rename = "includePricing")]
|
||||||
|
pub include_pricing: Option<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(FromRow)]
|
||||||
|
struct AvailRow {
|
||||||
|
account_id: i64,
|
||||||
|
model_name: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(FromRow)]
|
||||||
|
struct AccountRow {
|
||||||
|
id: i64,
|
||||||
|
site_id: i64,
|
||||||
|
username: Option<String>,
|
||||||
|
balance: Option<f64>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(FromRow)]
|
||||||
|
struct TokenRow {
|
||||||
|
id: i64,
|
||||||
|
account_id: i64,
|
||||||
|
name: String,
|
||||||
|
is_default: Option<bool>,
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn get_marketplace(
|
||||||
|
State(state): State<AppState>,
|
||||||
|
Query(_query): Query<MarketplaceQuery>,
|
||||||
|
) -> AppResult<Json<Value>> {
|
||||||
|
let value = marketplace_inner(&state.db).await?;
|
||||||
|
Ok(Json(value))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Pool-only entry point. Identical to `get_marketplace` but decoupled from
|
||||||
|
/// `AppState` and axum extractors so unit tests can drive it with an
|
||||||
|
/// in-memory SQLite pool (mirrors the `rebuild_routes_inner` pattern in
|
||||||
|
/// `services::route_rebuild`).
|
||||||
|
pub(crate) async fn marketplace_inner(
|
||||||
|
pool: &sqlx::Pool<sqlx::Sqlite>,
|
||||||
|
) -> AppResult<Value> {
|
||||||
|
// 1. Active (account, model) pairs. Filtering on accounts.status =
|
||||||
|
// 'active' here keeps the rest of the aggregation small.
|
||||||
|
let avail_rows: Vec<AvailRow> = sqlx::query_as(
|
||||||
|
"SELECT ma.account_id, ma.model_name
|
||||||
|
FROM model_availability ma
|
||||||
|
JOIN accounts a ON a.id = ma.account_id
|
||||||
|
WHERE COALESCE(ma.available, 0) = 1
|
||||||
|
AND a.status = 'active'",
|
||||||
|
)
|
||||||
|
.fetch_all(pool)
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
if avail_rows.is_empty() {
|
||||||
|
return Ok(json!({ "models": [], "meta": {} }));
|
||||||
|
}
|
||||||
|
|
||||||
|
// 2. One-shot lookups — accounts, sites, tokens. Bounded by the active
|
||||||
|
// account count (typically dozens), so a full scan beats per-model
|
||||||
|
// in-loop selects.
|
||||||
|
let accounts: Vec<AccountRow> = sqlx::query_as(
|
||||||
|
"SELECT id, site_id, username, balance FROM accounts WHERE status = 'active'",
|
||||||
|
)
|
||||||
|
.fetch_all(pool)
|
||||||
|
.await?;
|
||||||
|
let accounts_by_id: HashMap<i64, AccountRow> =
|
||||||
|
accounts.into_iter().map(|a| (a.id, a)).collect();
|
||||||
|
|
||||||
|
let sites: Vec<(i64, String)> = sqlx::query_as("SELECT id, name FROM sites")
|
||||||
|
.fetch_all(pool)
|
||||||
|
.await?;
|
||||||
|
let site_name_by_id: HashMap<i64, String> = sites.into_iter().collect();
|
||||||
|
|
||||||
|
let tokens: Vec<TokenRow> = sqlx::query_as(
|
||||||
|
"SELECT id, account_id, name, is_default
|
||||||
|
FROM account_tokens
|
||||||
|
WHERE COALESCE(enabled, 1) = 1",
|
||||||
|
)
|
||||||
|
.fetch_all(pool)
|
||||||
|
.await?;
|
||||||
|
let mut tokens_by_account: HashMap<i64, Vec<TokenRow>> = HashMap::new();
|
||||||
|
for t in tokens {
|
||||||
|
tokens_by_account.entry(t.account_id).or_default().push(t);
|
||||||
|
}
|
||||||
|
|
||||||
|
// 3. Group by model name. BTreeMap → output sorted by name (matches the
|
||||||
|
// deterministic ordering the frontend's table view expects).
|
||||||
|
let mut by_model: BTreeMap<String, Vec<i64>> = BTreeMap::new();
|
||||||
|
for row in avail_rows {
|
||||||
|
by_model
|
||||||
|
.entry(row.model_name)
|
||||||
|
.or_default()
|
||||||
|
.push(row.account_id);
|
||||||
|
}
|
||||||
|
|
||||||
|
// 4. Render. Stage 1 returns null/empty for the metadata fields; the
|
||||||
|
// frontend has fallback handling for each one.
|
||||||
|
let mut models: Vec<Value> = Vec::with_capacity(by_model.len());
|
||||||
|
for (model_name, account_ids) in by_model {
|
||||||
|
let account_count = account_ids.len();
|
||||||
|
let mut token_count: usize = 0;
|
||||||
|
let mut accounts_json: Vec<Value> = Vec::with_capacity(account_count);
|
||||||
|
for aid in &account_ids {
|
||||||
|
let acc = match accounts_by_id.get(aid) {
|
||||||
|
Some(a) => a,
|
||||||
|
None => continue,
|
||||||
|
};
|
||||||
|
let site_name = site_name_by_id
|
||||||
|
.get(&acc.site_id)
|
||||||
|
.cloned()
|
||||||
|
.unwrap_or_default();
|
||||||
|
let acc_tokens: &[TokenRow] = tokens_by_account
|
||||||
|
.get(aid)
|
||||||
|
.map(Vec::as_slice)
|
||||||
|
.unwrap_or(&[]);
|
||||||
|
token_count += acc_tokens.len();
|
||||||
|
let token_json: Vec<Value> = acc_tokens
|
||||||
|
.iter()
|
||||||
|
.map(|t| {
|
||||||
|
json!({
|
||||||
|
"id": t.id,
|
||||||
|
"name": t.name,
|
||||||
|
"isDefault": t.is_default.unwrap_or(false),
|
||||||
|
})
|
||||||
|
})
|
||||||
|
.collect();
|
||||||
|
accounts_json.push(json!({
|
||||||
|
"id": acc.id,
|
||||||
|
"site": site_name,
|
||||||
|
"username": acc.username,
|
||||||
|
"latency": Value::Null,
|
||||||
|
"balance": acc.balance.unwrap_or(0.0),
|
||||||
|
"tokens": token_json,
|
||||||
|
}));
|
||||||
|
}
|
||||||
|
|
||||||
|
models.push(json!({
|
||||||
|
"name": model_name,
|
||||||
|
"accountCount": account_count,
|
||||||
|
"tokenCount": token_count,
|
||||||
|
"avgLatency": Value::Null,
|
||||||
|
"successRate": Value::Null,
|
||||||
|
"description": Value::Null,
|
||||||
|
"tags": [],
|
||||||
|
"supportedEndpointTypes": [],
|
||||||
|
"pricingSources": [],
|
||||||
|
"accounts": accounts_json,
|
||||||
|
}));
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(json!({
|
||||||
|
"models": models,
|
||||||
|
"meta": {},
|
||||||
|
}))
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
use sqlx::sqlite::SqlitePoolOptions;
|
||||||
|
use sqlx::{Pool, Sqlite};
|
||||||
|
|
||||||
|
/// Minimal schema for the marketplace aggregation. Mirrors
|
||||||
|
/// `services::route_rebuild::tests::fixture_pool` — we create only the
|
||||||
|
/// columns `marketplace_inner` actually reads, so adding unrelated
|
||||||
|
/// migrations later doesn't break these tests.
|
||||||
|
async fn fixture_pool() -> Pool<Sqlite> {
|
||||||
|
let pool = SqlitePoolOptions::new()
|
||||||
|
.max_connections(1)
|
||||||
|
.connect("sqlite::memory:")
|
||||||
|
.await
|
||||||
|
.expect("open in-memory sqlite");
|
||||||
|
|
||||||
|
for ddl in [
|
||||||
|
"CREATE TABLE sites (
|
||||||
|
id INTEGER PRIMARY KEY,
|
||||||
|
name TEXT NOT NULL
|
||||||
|
)",
|
||||||
|
"CREATE TABLE accounts (
|
||||||
|
id INTEGER PRIMARY KEY,
|
||||||
|
site_id INTEGER NOT NULL,
|
||||||
|
username TEXT,
|
||||||
|
balance REAL,
|
||||||
|
status TEXT NOT NULL DEFAULT 'active'
|
||||||
|
)",
|
||||||
|
"CREATE TABLE model_availability (
|
||||||
|
account_id INTEGER NOT NULL,
|
||||||
|
model_name TEXT NOT NULL,
|
||||||
|
available INTEGER NOT NULL DEFAULT 1,
|
||||||
|
PRIMARY KEY (account_id, model_name)
|
||||||
|
)",
|
||||||
|
"CREATE TABLE account_tokens (
|
||||||
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||||
|
account_id INTEGER NOT NULL,
|
||||||
|
name TEXT NOT NULL,
|
||||||
|
enabled INTEGER DEFAULT 1,
|
||||||
|
is_default INTEGER DEFAULT 0
|
||||||
|
)",
|
||||||
|
] {
|
||||||
|
sqlx::query(ddl)
|
||||||
|
.execute(&pool)
|
||||||
|
.await
|
||||||
|
.unwrap_or_else(|err| panic!("failed to create schema: {err}\nddl: {ddl}"));
|
||||||
|
}
|
||||||
|
|
||||||
|
pool
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Empty DB → empty `models` array with empty `meta`. Locks the early-return
|
||||||
|
/// path at the top of `marketplace_inner` so an accidental refactor that
|
||||||
|
/// drops the guard still produces a frontend-safe response shape rather
|
||||||
|
/// than 500-ing on the follow-up `accounts`/`sites` selects.
|
||||||
|
#[tokio::test]
|
||||||
|
async fn empty_db_returns_empty_models_array() {
|
||||||
|
let pool = fixture_pool().await;
|
||||||
|
let value = marketplace_inner(&pool).await.expect("inner ok");
|
||||||
|
assert_eq!(value, json!({ "models": [], "meta": {} }));
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Inactive accounts and unavailable models must be filtered out at the
|
||||||
|
/// SQL level. Account 2 is `disabled` and account 3's row has
|
||||||
|
/// `available = 0` — neither should contribute to the response.
|
||||||
|
#[tokio::test]
|
||||||
|
async fn filters_inactive_accounts_and_unavailable_models() {
|
||||||
|
let pool = fixture_pool().await;
|
||||||
|
|
||||||
|
sqlx::query("INSERT INTO sites (id, name) VALUES (1, 'site-a')")
|
||||||
|
.execute(&pool).await.unwrap();
|
||||||
|
sqlx::query(
|
||||||
|
"INSERT INTO accounts (id, site_id, username, balance, status) VALUES
|
||||||
|
(1, 1, 'alice', 10.0, 'active'),
|
||||||
|
(2, 1, 'bob', 5.0, 'disabled'),
|
||||||
|
(3, 1, 'carol', 7.0, 'active')",
|
||||||
|
).execute(&pool).await.unwrap();
|
||||||
|
sqlx::query(
|
||||||
|
"INSERT INTO model_availability (account_id, model_name, available) VALUES
|
||||||
|
(1, 'gpt-x', 1),
|
||||||
|
(2, 'gpt-x', 1),
|
||||||
|
(3, 'gpt-x', 0)",
|
||||||
|
).execute(&pool).await.unwrap();
|
||||||
|
|
||||||
|
let value = marketplace_inner(&pool).await.expect("inner ok");
|
||||||
|
let models = value.get("models").and_then(Value::as_array).unwrap();
|
||||||
|
assert_eq!(models.len(), 1, "exactly one model should survive filtering");
|
||||||
|
let m = &models[0];
|
||||||
|
assert_eq!(m["name"], "gpt-x");
|
||||||
|
assert_eq!(m["accountCount"], 1, "only account 1 (active + available) counts");
|
||||||
|
let accounts = m["accounts"].as_array().unwrap();
|
||||||
|
assert_eq!(accounts.len(), 1);
|
||||||
|
assert_eq!(accounts[0]["username"], "alice");
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Multi-account, multi-token aggregation. Verifies:
|
||||||
|
/// - `accountCount` == distinct active accounts advertising the model
|
||||||
|
/// - `tokenCount` == sum of enabled tokens across those accounts
|
||||||
|
/// - per-account `tokens[]` is correctly nested with `isDefault`
|
||||||
|
/// - disabled tokens (`enabled = 0`) don't show up
|
||||||
|
/// - site name is resolved from `sites.name`
|
||||||
|
#[tokio::test]
|
||||||
|
async fn aggregates_accounts_and_tokens_per_model() {
|
||||||
|
let pool = fixture_pool().await;
|
||||||
|
|
||||||
|
sqlx::query("INSERT INTO sites (id, name) VALUES (1, 'openai-mirror'), (2, 'anthropic-mirror')")
|
||||||
|
.execute(&pool).await.unwrap();
|
||||||
|
sqlx::query(
|
||||||
|
"INSERT INTO accounts (id, site_id, username, balance, status) VALUES
|
||||||
|
(1, 1, 'alice', 100.5, 'active'),
|
||||||
|
(2, 2, 'bob', 50.0, 'active')",
|
||||||
|
).execute(&pool).await.unwrap();
|
||||||
|
sqlx::query(
|
||||||
|
"INSERT INTO model_availability (account_id, model_name, available) VALUES
|
||||||
|
(1, 'gpt-x', 1),
|
||||||
|
(2, 'gpt-x', 1)",
|
||||||
|
).execute(&pool).await.unwrap();
|
||||||
|
// Account 1: two enabled tokens (one default), one disabled token (must
|
||||||
|
// be excluded). Account 2: one enabled non-default token.
|
||||||
|
sqlx::query(
|
||||||
|
"INSERT INTO account_tokens (id, account_id, name, enabled, is_default) VALUES
|
||||||
|
(10, 1, 'primary', 1, 1),
|
||||||
|
(11, 1, 'backup', 1, 0),
|
||||||
|
(12, 1, 'retired', 0, 0),
|
||||||
|
(20, 2, 'only-one', 1, 0)",
|
||||||
|
).execute(&pool).await.unwrap();
|
||||||
|
|
||||||
|
let value = marketplace_inner(&pool).await.expect("inner ok");
|
||||||
|
let models = value["models"].as_array().unwrap();
|
||||||
|
assert_eq!(models.len(), 1);
|
||||||
|
let m = &models[0];
|
||||||
|
assert_eq!(m["name"], "gpt-x");
|
||||||
|
assert_eq!(m["accountCount"], 2);
|
||||||
|
assert_eq!(m["tokenCount"], 3, "2 from account 1 + 1 from account 2; disabled token excluded");
|
||||||
|
// Stage 1 placeholders.
|
||||||
|
assert!(m["avgLatency"].is_null());
|
||||||
|
assert!(m["successRate"].is_null());
|
||||||
|
assert!(m["description"].is_null());
|
||||||
|
assert_eq!(m["tags"], json!([]));
|
||||||
|
assert_eq!(m["supportedEndpointTypes"], json!([]));
|
||||||
|
assert_eq!(m["pricingSources"], json!([]));
|
||||||
|
|
||||||
|
let accounts = m["accounts"].as_array().unwrap();
|
||||||
|
assert_eq!(accounts.len(), 2);
|
||||||
|
|
||||||
|
// Find each account by id rather than relying on insertion order in
|
||||||
|
// the HashMap-driven account_ids vec.
|
||||||
|
let alice = accounts.iter().find(|a| a["id"] == 1).expect("alice present");
|
||||||
|
assert_eq!(alice["site"], "openai-mirror");
|
||||||
|
assert_eq!(alice["username"], "alice");
|
||||||
|
assert_eq!(alice["balance"], 100.5);
|
||||||
|
let alice_tokens = alice["tokens"].as_array().unwrap();
|
||||||
|
assert_eq!(alice_tokens.len(), 2, "disabled 'retired' token must not be exposed");
|
||||||
|
let primary = alice_tokens.iter().find(|t| t["id"] == 10).unwrap();
|
||||||
|
assert_eq!(primary["name"], "primary");
|
||||||
|
assert_eq!(primary["isDefault"], true);
|
||||||
|
let backup = alice_tokens.iter().find(|t| t["id"] == 11).unwrap();
|
||||||
|
assert_eq!(backup["isDefault"], false);
|
||||||
|
|
||||||
|
let bob = accounts.iter().find(|a| a["id"] == 2).expect("bob present");
|
||||||
|
assert_eq!(bob["site"], "anthropic-mirror");
|
||||||
|
assert_eq!(bob["balance"], 50.0);
|
||||||
|
assert_eq!(bob["tokens"].as_array().unwrap().len(), 1);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Models must come back in deterministic name order (BTreeMap). The
|
||||||
|
/// frontend's table view assumes a stable sort so toggling refresh
|
||||||
|
/// doesn't reshuffle rows; lock that contract here.
|
||||||
|
#[tokio::test]
|
||||||
|
async fn models_are_sorted_by_name() {
|
||||||
|
let pool = fixture_pool().await;
|
||||||
|
|
||||||
|
sqlx::query("INSERT INTO sites (id, name) VALUES (1, 's')")
|
||||||
|
.execute(&pool).await.unwrap();
|
||||||
|
sqlx::query("INSERT INTO accounts (id, site_id, status) VALUES (1, 1, 'active')")
|
||||||
|
.execute(&pool).await.unwrap();
|
||||||
|
// Insert in scrambled order; expect output sorted lexicographically.
|
||||||
|
sqlx::query(
|
||||||
|
"INSERT INTO model_availability (account_id, model_name, available) VALUES
|
||||||
|
(1, 'gpt-z', 1),
|
||||||
|
(1, 'gpt-a', 1),
|
||||||
|
(1, 'gpt-m', 1)",
|
||||||
|
).execute(&pool).await.unwrap();
|
||||||
|
|
||||||
|
let value = marketplace_inner(&pool).await.expect("inner ok");
|
||||||
|
let names: Vec<&str> = value["models"]
|
||||||
|
.as_array()
|
||||||
|
.unwrap()
|
||||||
|
.iter()
|
||||||
|
.map(|m| m["name"].as_str().unwrap())
|
||||||
|
.collect();
|
||||||
|
assert_eq!(names, vec!["gpt-a", "gpt-m", "gpt-z"]);
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -356,6 +356,82 @@ async fn site_api_endpoints(state: &AppState, site_id: i64) -> AppResult<Vec<Val
|
|||||||
.collect())
|
.collect())
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Deserialize)]
|
||||||
|
pub struct UpdateDisabledModelsRequest {
|
||||||
|
#[serde(default)]
|
||||||
|
pub models: Vec<String>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// GET /api/sites/{id}/disabled-models — list disabled model names for a site.
|
||||||
|
pub async fn list_site_disabled_models(
|
||||||
|
State(state): State<AppState>,
|
||||||
|
Path(id): Path<i64>,
|
||||||
|
) -> AppResult<Json<Value>> {
|
||||||
|
// Validate the site exists so we return a clean 404 instead of an empty list
|
||||||
|
// for unknown ids.
|
||||||
|
let exists: Option<i64> =
|
||||||
|
sqlx::query_scalar("SELECT id FROM sites WHERE id = ?")
|
||||||
|
.bind(id)
|
||||||
|
.fetch_optional(&*state.db)
|
||||||
|
.await?;
|
||||||
|
if exists.is_none() {
|
||||||
|
return Err(AppError::NotFound("site not found".into()));
|
||||||
|
}
|
||||||
|
|
||||||
|
let models: Vec<String> =
|
||||||
|
sqlx::query_scalar("SELECT model_name FROM site_disabled_models WHERE site_id = ? ORDER BY model_name")
|
||||||
|
.bind(id)
|
||||||
|
.fetch_all(&*state.db)
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
Ok(Json(json!({ "models": models })))
|
||||||
|
}
|
||||||
|
|
||||||
|
/// PUT /api/sites/{id}/disabled-models — replace the disabled-model list for a site.
|
||||||
|
pub async fn update_site_disabled_models(
|
||||||
|
State(state): State<AppState>,
|
||||||
|
Path(id): Path<i64>,
|
||||||
|
Json(payload): Json<UpdateDisabledModelsRequest>,
|
||||||
|
) -> AppResult<Json<Value>> {
|
||||||
|
let exists: Option<i64> =
|
||||||
|
sqlx::query_scalar("SELECT id FROM sites WHERE id = ?")
|
||||||
|
.bind(id)
|
||||||
|
.fetch_optional(&*state.db)
|
||||||
|
.await?;
|
||||||
|
if exists.is_none() {
|
||||||
|
return Err(AppError::NotFound("site not found".into()));
|
||||||
|
}
|
||||||
|
|
||||||
|
// Normalise + dedupe the inbound list. Empty strings are dropped.
|
||||||
|
let mut normalised: Vec<String> = payload
|
||||||
|
.models
|
||||||
|
.into_iter()
|
||||||
|
.map(|m| m.trim().to_string())
|
||||||
|
.filter(|m| !m.is_empty())
|
||||||
|
.collect();
|
||||||
|
normalised.sort();
|
||||||
|
normalised.dedup();
|
||||||
|
|
||||||
|
let mut tx = state.db.begin().await?;
|
||||||
|
sqlx::query("DELETE FROM site_disabled_models WHERE site_id = ?")
|
||||||
|
.bind(id)
|
||||||
|
.execute(&mut *tx)
|
||||||
|
.await?;
|
||||||
|
for model in &normalised {
|
||||||
|
sqlx::query(
|
||||||
|
"INSERT INTO site_disabled_models (site_id, model_name) VALUES (?, ?)\n ON CONFLICT(site_id, model_name) DO NOTHING",
|
||||||
|
)
|
||||||
|
.bind(id)
|
||||||
|
.bind(model)
|
||||||
|
.execute(&mut *tx)
|
||||||
|
.await?;
|
||||||
|
}
|
||||||
|
tx.commit().await?;
|
||||||
|
|
||||||
|
info!(site_id = id, count = normalised.len(), "updated site disabled models");
|
||||||
|
Ok(Json(json!({ "models": normalised })))
|
||||||
|
}
|
||||||
|
|
||||||
async fn replace_site_api_endpoints(
|
async fn replace_site_api_endpoints(
|
||||||
tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>,
|
tx: &mut sqlx::Transaction<'_, sqlx::Sqlite>,
|
||||||
site_id: i64,
|
site_id: i64,
|
||||||
|
|||||||
+78
-1
@@ -3,7 +3,7 @@
|
|||||||
|
|
||||||
use crate::{
|
use crate::{
|
||||||
auth::{jwt::generate_jwt, password::verify_password},
|
auth::{jwt::generate_jwt, password::verify_password},
|
||||||
db::models::user::User,
|
db::models::{downstream_api_key::DownstreamApiKey, user::User},
|
||||||
error::{AppError, AppResult},
|
error::{AppError, AppResult},
|
||||||
AppState,
|
AppState,
|
||||||
};
|
};
|
||||||
@@ -32,6 +32,12 @@ pub struct LoginRequest {
|
|||||||
pub password: String,
|
pub password: String,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Deserialize)]
|
||||||
|
pub struct LoginApiKeyRequest {
|
||||||
|
#[serde(alias = "apiKey", alias = "key")]
|
||||||
|
pub api_key: String,
|
||||||
|
}
|
||||||
|
|
||||||
#[allow(dead_code)]
|
#[allow(dead_code)]
|
||||||
#[derive(Debug, Serialize)]
|
#[derive(Debug, Serialize)]
|
||||||
pub struct AuthResponse {
|
pub struct AuthResponse {
|
||||||
@@ -56,6 +62,8 @@ pub fn routes() -> Router<AppState> {
|
|||||||
.route("/api/users/register", post(register))
|
.route("/api/users/register", post(register))
|
||||||
.route("/api/auth/login", post(login))
|
.route("/api/auth/login", post(login))
|
||||||
.route("/api/users/login", post(login))
|
.route("/api/users/login", post(login))
|
||||||
|
.route("/api/auth/login-api-key", post(login_api_key))
|
||||||
|
.route("/api/users/login-api-key", post(login_api_key))
|
||||||
.route("/api/auth/me", get(me))
|
.route("/api/auth/me", get(me))
|
||||||
.route("/api/users/me", get(me).patch(update_me))
|
.route("/api/users/me", get(me).patch(update_me))
|
||||||
.route("/api/users/me/password", post(change_password))
|
.route("/api/users/me/password", post(change_password))
|
||||||
@@ -197,6 +205,75 @@ async fn login(
|
|||||||
})))
|
})))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// POST /api/users/login-api-key — exchange a downstream API key for a user JWT.
|
||||||
|
///
|
||||||
|
/// The API key must belong to a managed user (downstream_api_keys.user_id IS NOT NULL),
|
||||||
|
/// be enabled, and the owning user must be active.
|
||||||
|
async fn login_api_key(
|
||||||
|
State(state): State<AppState>,
|
||||||
|
Json(payload): Json<LoginApiKeyRequest>,
|
||||||
|
) -> AppResult<Json<Value>> {
|
||||||
|
let key = payload.api_key.trim();
|
||||||
|
if key.is_empty() {
|
||||||
|
return Err(AppError::BadRequest("api_key is required".into()));
|
||||||
|
}
|
||||||
|
|
||||||
|
let row: Option<DownstreamApiKey> = sqlx::query_as::<_, DownstreamApiKey>(
|
||||||
|
"SELECT * FROM downstream_api_keys WHERE key = ? AND COALESCE(enabled, 1) = 1",
|
||||||
|
)
|
||||||
|
.bind(key)
|
||||||
|
.fetch_optional(&*state.db)
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
let api_key = row.ok_or_else(|| AppError::Unauthorized("Invalid API key".into()))?;
|
||||||
|
let user_id = api_key
|
||||||
|
.user_id
|
||||||
|
.ok_or_else(|| AppError::Unauthorized("API key is not bound to a user".into()))?;
|
||||||
|
|
||||||
|
let user: User = sqlx::query_as::<_, User>("SELECT * FROM users WHERE id = ?")
|
||||||
|
.bind(user_id)
|
||||||
|
.fetch_optional(&*state.db)
|
||||||
|
.await?
|
||||||
|
.ok_or_else(|| AppError::Unauthorized("User not found".into()))?;
|
||||||
|
|
||||||
|
if user.status != "active" {
|
||||||
|
return Err(AppError::Forbidden("Account is disabled".into()));
|
||||||
|
}
|
||||||
|
|
||||||
|
info!("Login via API key: user_id={}", user.id);
|
||||||
|
|
||||||
|
let token = generate_jwt(
|
||||||
|
&state.config.auth.jwt_secret,
|
||||||
|
&user.id.to_string(),
|
||||||
|
&user.role,
|
||||||
|
state.config.auth.jwt_expiry_hours,
|
||||||
|
)
|
||||||
|
.map_err(|e| AppError::Internal(anyhow::anyhow!("Failed to generate JWT: {e}")))?;
|
||||||
|
|
||||||
|
Ok(Json(json!({
|
||||||
|
"token": token,
|
||||||
|
"user": {
|
||||||
|
"id": user.id,
|
||||||
|
"username": user.username,
|
||||||
|
"email": user.email,
|
||||||
|
"role": user.role,
|
||||||
|
"status": user.status,
|
||||||
|
"createdAt": user.created_at,
|
||||||
|
},
|
||||||
|
"data": {
|
||||||
|
"token": token,
|
||||||
|
"user": {
|
||||||
|
"id": user.id,
|
||||||
|
"username": user.username,
|
||||||
|
"email": user.email,
|
||||||
|
"role": user.role,
|
||||||
|
"status": user.status,
|
||||||
|
"createdAt": user.created_at,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
})))
|
||||||
|
}
|
||||||
|
|
||||||
/// GET /api/auth/me — returns current user profile from JWT
|
/// GET /api/auth/me — returns current user profile from JWT
|
||||||
async fn me(State(state): State<AppState>, headers: HeaderMap) -> AppResult<Json<Value>> {
|
async fn me(State(state): State<AppState>, headers: HeaderMap) -> AppResult<Json<Value>> {
|
||||||
let auth = headers
|
let auth = headers
|
||||||
|
|||||||
@@ -0,0 +1,447 @@
|
|||||||
|
//! `/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.
|
||||||
|
//!
|
||||||
|
//! NO path rewriting on the response — ComfyUI's HTML/JS reference
|
||||||
|
//! root-relative URLs which the browser will request under `/comfy/...`,
|
||||||
|
//! and the path stripping done here lines them up with the upstream.
|
||||||
|
|
||||||
|
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) {
|
||||||
|
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")
|
||||||
|
}
|
||||||
|
|
||||||
|
// ── 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()
|
||||||
|
}
|
||||||
+331
-15
@@ -1,4 +1,4 @@
|
|||||||
use crate::AppState;
|
use crate::{auth::jwt::validate_jwt, AppState};
|
||||||
use axum::{
|
use axum::{
|
||||||
extract::State,
|
extract::State,
|
||||||
http::HeaderMap,
|
http::HeaderMap,
|
||||||
@@ -62,12 +62,14 @@ emotion: 平静
|
|||||||
pub fn routes() -> Router<AppState> {
|
pub fn routes() -> Router<AppState> {
|
||||||
Router::new()
|
Router::new()
|
||||||
.route("/api/comfyui-agent/status", get(status_handler))
|
.route("/api/comfyui-agent/status", get(status_handler))
|
||||||
|
.route("/api/comfyui-agent/models", get(models_handler))
|
||||||
.route("/api/comfyui-agent/chat", post(chat_handler))
|
.route("/api/comfyui-agent/chat", post(chat_handler))
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Debug, Deserialize)]
|
#[derive(Debug, Deserialize)]
|
||||||
struct AgentChatRequest {
|
struct AgentChatRequest {
|
||||||
messages: Vec<AgentMessage>,
|
messages: Vec<AgentMessage>,
|
||||||
|
model: Option<String>,
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Debug, Deserialize)]
|
#[derive(Debug, Deserialize)]
|
||||||
@@ -76,18 +78,78 @@ struct AgentMessage {
|
|||||||
content: String,
|
content: String,
|
||||||
}
|
}
|
||||||
|
|
||||||
fn upstream_config(state: &AppState) -> (String, String) {
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||||
let base = std::env::var("BOOSAPI_UPSTREAM_BASE")
|
enum AgentMode {
|
||||||
|
/// Route through this server's own `/v1/*` using the built-in `comfyui-agent` downstream key.
|
||||||
|
Local,
|
||||||
|
/// Legacy fallback: external upstream configured via `BOOSAPI_UPSTREAM_BASE` / `BOOSAPI_UPSTREAM_KEY`.
|
||||||
|
FallbackRemote,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl AgentMode {
|
||||||
|
fn as_str(self) -> &'static str {
|
||||||
|
match self {
|
||||||
|
AgentMode::Local => "local",
|
||||||
|
AgentMode::FallbackRemote => "fallback_remote",
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn detect_agent_mode() -> AgentMode {
|
||||||
|
let base_override = std::env::var("BOOSAPI_UPSTREAM_BASE")
|
||||||
.ok()
|
.ok()
|
||||||
.map(|value| value.trim().trim_end_matches('/').to_string())
|
.map(|v| v.trim().to_string())
|
||||||
.filter(|value| !value.is_empty())
|
.filter(|v| !v.is_empty());
|
||||||
.unwrap_or_else(|| "https://api.boos.lat/v1".to_string());
|
let key_override = std::env::var("BOOSAPI_UPSTREAM_KEY")
|
||||||
let key = std::env::var("BOOSAPI_UPSTREAM_KEY")
|
|
||||||
.ok()
|
.ok()
|
||||||
.map(|value| value.trim().to_string())
|
.map(|v| v.trim().to_string())
|
||||||
.filter(|value| !value.is_empty())
|
.filter(|v| !v.is_empty());
|
||||||
.unwrap_or_else(|| state.config.auth.proxy_token.clone());
|
if base_override.is_some() || key_override.is_some() {
|
||||||
(base, key)
|
AgentMode::FallbackRemote
|
||||||
|
} else {
|
||||||
|
AgentMode::Local
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn local_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)
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Resolve `(mode, api_base, api_key)` for the agent's outbound calls.
|
||||||
|
///
|
||||||
|
/// Defaults to `AgentMode::Local`, which targets this server's own `/v1`
|
||||||
|
/// endpoints with the built-in `comfyui-agent` downstream key — every call
|
||||||
|
/// then runs through the regular accounting / routing pipeline.
|
||||||
|
///
|
||||||
|
/// Setting `BOOSAPI_UPSTREAM_BASE` or `BOOSAPI_UPSTREAM_KEY` switches to
|
||||||
|
/// `FallbackRemote` for operators that still want to point straight at an
|
||||||
|
/// external aggregator.
|
||||||
|
fn upstream_config(state: &AppState) -> (AgentMode, String, String) {
|
||||||
|
match detect_agent_mode() {
|
||||||
|
AgentMode::Local => {
|
||||||
|
let base = local_base(state);
|
||||||
|
let key = state.builtin_keys.comfyui_agent.1.clone();
|
||||||
|
(AgentMode::Local, base, key)
|
||||||
|
}
|
||||||
|
AgentMode::FallbackRemote => {
|
||||||
|
let base = std::env::var("BOOSAPI_UPSTREAM_BASE")
|
||||||
|
.ok()
|
||||||
|
.map(|value| value.trim().trim_end_matches('/').to_string())
|
||||||
|
.filter(|value| !value.is_empty())
|
||||||
|
.unwrap_or_else(|| "https://api.boos.lat/v1".to_string());
|
||||||
|
let key = std::env::var("BOOSAPI_UPSTREAM_KEY")
|
||||||
|
.ok()
|
||||||
|
.map(|value| value.trim().to_string())
|
||||||
|
.filter(|value| !value.is_empty())
|
||||||
|
.unwrap_or_else(|| state.config.auth.proxy_token.clone());
|
||||||
|
(AgentMode::FallbackRemote, base, key)
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
fn authorize(headers: &HeaderMap, state: &AppState) -> bool {
|
fn authorize(headers: &HeaderMap, state: &AppState) -> bool {
|
||||||
@@ -101,7 +163,21 @@ fn authorize(headers: &HeaderMap, state: &AppState) -> bool {
|
|||||||
.or_else(|| auth_header.strip_prefix("bearer "))
|
.or_else(|| auth_header.strip_prefix("bearer "))
|
||||||
.unwrap_or_default();
|
.unwrap_or_default();
|
||||||
|
|
||||||
token == state.config.auth.auth_token
|
if token == state.config.auth.auth_token {
|
||||||
|
return true;
|
||||||
|
}
|
||||||
|
|
||||||
|
validate_jwt(&state.config.auth.jwt_secret, token)
|
||||||
|
.map(|claims| claims.role == "admin" || claims.role == "user")
|
||||||
|
.unwrap_or(false)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn is_exact_model_pattern(value: &str) -> bool {
|
||||||
|
let normalized = value.trim();
|
||||||
|
!normalized.is_empty()
|
||||||
|
&& !normalized.to_ascii_lowercase().starts_with("re:")
|
||||||
|
&& !normalized.contains('*')
|
||||||
|
&& !normalized.contains('?')
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn status_handler(
|
async fn status_handler(
|
||||||
@@ -114,15 +190,62 @@ async fn status_handler(
|
|||||||
Json(json!({ "error": "Unauthorized" })),
|
Json(json!({ "error": "Unauthorized" })),
|
||||||
));
|
));
|
||||||
}
|
}
|
||||||
let (upstream_base, _) = upstream_config(&state);
|
let (mode, upstream_base, _) = upstream_config(&state);
|
||||||
|
let key_name = match mode {
|
||||||
|
AgentMode::Local => "comfyui-agent",
|
||||||
|
AgentMode::FallbackRemote => "external",
|
||||||
|
};
|
||||||
Ok(Json(json!({
|
Ok(Json(json!({
|
||||||
"success": true,
|
"success": true,
|
||||||
"model": state.config.comfyui.agent_model,
|
"model": state.config.comfyui.agent_model,
|
||||||
|
"mode": mode.as_str(),
|
||||||
|
"keyName": key_name,
|
||||||
"upstreamBase": upstream_base,
|
"upstreamBase": upstream_base,
|
||||||
"comfyui": state.comfyui.status_response().await,
|
"comfyui": state.comfyui.status_response().await,
|
||||||
})))
|
})))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
async fn models_handler(
|
||||||
|
State(state): State<AppState>,
|
||||||
|
headers: HeaderMap,
|
||||||
|
) -> Result<Json<Value>, (axum::http::StatusCode, Json<Value>)> {
|
||||||
|
if !authorize(&headers, &state) {
|
||||||
|
return Err((
|
||||||
|
axum::http::StatusCode::UNAUTHORIZED,
|
||||||
|
Json(json!({ "error": "Unauthorized" })),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
|
||||||
|
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_exact_model_pattern(value))
|
||||||
|
.filter(|value| seen.insert(value.clone()))
|
||||||
|
.collect();
|
||||||
|
|
||||||
|
Ok(Json(json!({
|
||||||
|
"success": true,
|
||||||
|
"models": models,
|
||||||
|
"defaultModel": state.config.comfyui.agent_model,
|
||||||
|
})))
|
||||||
|
}
|
||||||
|
|
||||||
#[derive(Debug, Default)]
|
#[derive(Debug, Default)]
|
||||||
struct ParsedSections {
|
struct ParsedSections {
|
||||||
characters: Vec<CharacterSection>,
|
characters: Vec<CharacterSection>,
|
||||||
@@ -373,6 +496,13 @@ async fn chat_handler(
|
|||||||
));
|
));
|
||||||
}
|
}
|
||||||
|
|
||||||
|
let requested_model = body
|
||||||
|
.model
|
||||||
|
.as_deref()
|
||||||
|
.map(str::trim)
|
||||||
|
.filter(|value| !value.is_empty())
|
||||||
|
.map(ToOwned::to_owned);
|
||||||
|
|
||||||
let messages: Vec<Value> = body
|
let messages: Vec<Value> = body
|
||||||
.messages
|
.messages
|
||||||
.into_iter()
|
.into_iter()
|
||||||
@@ -383,8 +513,13 @@ async fn chat_handler(
|
|||||||
})
|
})
|
||||||
})
|
})
|
||||||
.collect();
|
.collect();
|
||||||
let (api_base, api_key) = upstream_config(&state);
|
let (mode, api_base, api_key) = upstream_config(&state);
|
||||||
let agent_model = state.config.comfyui.agent_model.clone();
|
info!(
|
||||||
|
mode = mode.as_str(),
|
||||||
|
upstream_base = %api_base,
|
||||||
|
"comfyui-agent: resolved upstream",
|
||||||
|
);
|
||||||
|
let agent_model = requested_model.unwrap_or_else(|| state.config.comfyui.agent_model.clone());
|
||||||
let (tx, rx) = unbounded_channel::<Result<Event, Infallible>>();
|
let (tx, rx) = unbounded_channel::<Result<Event, Infallible>>();
|
||||||
|
|
||||||
tokio::spawn(async move {
|
tokio::spawn(async move {
|
||||||
@@ -723,3 +858,184 @@ async fn chat_handler(
|
|||||||
|
|
||||||
Ok(Sse::new(UnboundedReceiverStream::new(rx)))
|
Ok(Sse::new(UnboundedReceiverStream::new(rx)))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
use crate::config::{AppConfig, AuthConfig, DbConfig, ProxyConfig, ServerConfig};
|
||||||
|
use crate::services::builtin_keys::BuiltinKeys;
|
||||||
|
use axum::http::HeaderValue;
|
||||||
|
use sqlx::sqlite::SqlitePoolOptions;
|
||||||
|
use tokio::sync::Mutex;
|
||||||
|
use wiremock::matchers::{header, method, path};
|
||||||
|
use wiremock::{Mock, MockServer, ResponseTemplate};
|
||||||
|
|
||||||
|
/// `BOOSAPI_UPSTREAM_*` is process-global. Serialize tests that read or
|
||||||
|
/// mutate it so cargo's parallel runner does not race their assertions.
|
||||||
|
/// Uses `tokio::sync::Mutex` so the guard can safely span `.await` points
|
||||||
|
/// (mock server setup, agent worker poll loop) without tripping
|
||||||
|
/// `clippy::await_holding_lock`.
|
||||||
|
static ENV_GUARD: Mutex<()> = Mutex::const_new(());
|
||||||
|
|
||||||
|
async fn make_state(host: &str, port: u16) -> AppState {
|
||||||
|
let pool = SqlitePoolOptions::new()
|
||||||
|
.max_connections(1)
|
||||||
|
.connect("sqlite::memory:")
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
AppState::new(
|
||||||
|
AppConfig {
|
||||||
|
server: ServerConfig {
|
||||||
|
host: host.to_string(),
|
||||||
|
port,
|
||||||
|
log_level: "info".into(),
|
||||||
|
cors_origins: vec![],
|
||||||
|
},
|
||||||
|
db: DbConfig {
|
||||||
|
url: "sqlite::memory:".into(),
|
||||||
|
max_connections: 1,
|
||||||
|
ssl: false,
|
||||||
|
runtime_db_url: None,
|
||||||
|
runtime_db_dialect: None,
|
||||||
|
runtime_db_ssl: None,
|
||||||
|
},
|
||||||
|
auth: AuthConfig {
|
||||||
|
jwt_secret: "secret".into(),
|
||||||
|
jwt_expiry_hours: 1,
|
||||||
|
auth_token: "auth".into(),
|
||||||
|
proxy_token: "proxy".into(),
|
||||||
|
},
|
||||||
|
proxy: ProxyConfig {
|
||||||
|
upstream_timeout_secs: 30,
|
||||||
|
first_byte_timeout_secs: 30,
|
||||||
|
max_retries: 3,
|
||||||
|
proxy_log_retention_days: 7,
|
||||||
|
},
|
||||||
|
comfyui: Default::default(),
|
||||||
|
},
|
||||||
|
pool,
|
||||||
|
BuiltinKeys {
|
||||||
|
comfyui_agent: (1, "sk-comfyui-test".into()),
|
||||||
|
},
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn clear_upstream_env() {
|
||||||
|
std::env::remove_var("BOOSAPI_UPSTREAM_BASE");
|
||||||
|
std::env::remove_var("BOOSAPI_UPSTREAM_KEY");
|
||||||
|
}
|
||||||
|
|
||||||
|
fn auth_headers() -> HeaderMap {
|
||||||
|
let mut headers = HeaderMap::new();
|
||||||
|
headers.insert("Authorization", HeaderValue::from_static("Bearer auth"));
|
||||||
|
headers
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn call_status(state: AppState) -> Value {
|
||||||
|
let result = status_handler(State(state), auth_headers()).await;
|
||||||
|
result.expect("status_handler returned error").0
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 3.5 — `/api/comfyui-agent/status` reflects the env state in one shot:
|
||||||
|
/// no env → local; whitespace/empty env → still local (treated as unset);
|
||||||
|
/// real env → fallback_remote with trimmed `upstreamBase`.
|
||||||
|
#[tokio::test]
|
||||||
|
async fn status_reflects_env_state_across_three_scenarios() {
|
||||||
|
let _g = ENV_GUARD.lock().await;
|
||||||
|
clear_upstream_env();
|
||||||
|
|
||||||
|
// Scenario A: env unset → local; 0.0.0.0 host normalizes to 127.0.0.1.
|
||||||
|
{
|
||||||
|
let state = make_state("0.0.0.0", 7860).await;
|
||||||
|
let body = call_status(state).await;
|
||||||
|
assert_eq!(body["mode"], "local");
|
||||||
|
assert_eq!(body["keyName"], "comfyui-agent");
|
||||||
|
assert_eq!(body["upstreamBase"], "http://127.0.0.1:7860/v1");
|
||||||
|
}
|
||||||
|
|
||||||
|
// Scenario B: env set to whitespace/empty → treated as unset → local.
|
||||||
|
std::env::set_var("BOOSAPI_UPSTREAM_BASE", " ");
|
||||||
|
std::env::set_var("BOOSAPI_UPSTREAM_KEY", "");
|
||||||
|
{
|
||||||
|
let state = make_state("127.0.0.1", 9000).await;
|
||||||
|
let body = call_status(state).await;
|
||||||
|
assert_eq!(body["mode"], "local");
|
||||||
|
assert_eq!(body["keyName"], "comfyui-agent");
|
||||||
|
assert_eq!(body["upstreamBase"], "http://127.0.0.1:9000/v1");
|
||||||
|
}
|
||||||
|
|
||||||
|
// Scenario C: env set → fallback_remote, base trimmed of trailing slash.
|
||||||
|
std::env::set_var("BOOSAPI_UPSTREAM_BASE", "https://api.example.com/v1/");
|
||||||
|
std::env::remove_var("BOOSAPI_UPSTREAM_KEY");
|
||||||
|
{
|
||||||
|
let state = make_state("0.0.0.0", 7860).await;
|
||||||
|
let body = call_status(state).await;
|
||||||
|
assert_eq!(body["mode"], "fallback_remote");
|
||||||
|
assert_eq!(body["keyName"], "external");
|
||||||
|
assert_eq!(body["upstreamBase"], "https://api.example.com/v1");
|
||||||
|
}
|
||||||
|
|
||||||
|
clear_upstream_env();
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 3.6 — In Local mode the agent's chat call must hit the local
|
||||||
|
/// `/v1/chat/completions` endpoint with the built-in `comfyui-agent`
|
||||||
|
/// downstream key as the bearer token.
|
||||||
|
#[tokio::test]
|
||||||
|
async fn agent_local_mode_routes_chat_to_local_upstream() {
|
||||||
|
let _g = ENV_GUARD.lock().await;
|
||||||
|
clear_upstream_env();
|
||||||
|
|
||||||
|
let mock = MockServer::start().await;
|
||||||
|
let sse = "data: {\"choices\":[{\"delta\":{\"content\":\"hi\"}}]}\n\ndata: [DONE]\n\n";
|
||||||
|
Mock::given(method("POST"))
|
||||||
|
.and(path("/v1/chat/completions"))
|
||||||
|
.and(header("authorization", "Bearer sk-comfyui-test"))
|
||||||
|
.respond_with(
|
||||||
|
ResponseTemplate::new(200)
|
||||||
|
.insert_header("content-type", "text/event-stream")
|
||||||
|
.set_body_string(sse),
|
||||||
|
)
|
||||||
|
.mount(&mock)
|
||||||
|
.await;
|
||||||
|
|
||||||
|
let addr = mock.address();
|
||||||
|
let host = addr.ip().to_string();
|
||||||
|
let port = addr.port();
|
||||||
|
|
||||||
|
let state = make_state(&host, port).await;
|
||||||
|
|
||||||
|
let result = chat_handler(
|
||||||
|
State(state),
|
||||||
|
auth_headers(),
|
||||||
|
Json(AgentChatRequest {
|
||||||
|
messages: vec![AgentMessage {
|
||||||
|
role: "user".to_string(),
|
||||||
|
content: "hello".to_string(),
|
||||||
|
}],
|
||||||
|
}),
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
assert!(result.is_ok(), "chat_handler should accept the request");
|
||||||
|
// The Sse stream is dropped here; the worker still runs the upstream
|
||||||
|
// POST before its streaming loop notices `tx.closed()`.
|
||||||
|
|
||||||
|
// Wait up to ~2s for the spawned worker's POST to land at wiremock.
|
||||||
|
let mut hit = false;
|
||||||
|
for _ in 0..40 {
|
||||||
|
tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
|
||||||
|
let received = mock.received_requests().await.unwrap_or_default();
|
||||||
|
if received
|
||||||
|
.iter()
|
||||||
|
.any(|r| r.url.path() == "/v1/chat/completions")
|
||||||
|
{
|
||||||
|
hit = true;
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
}
|
||||||
|
assert!(
|
||||||
|
hit,
|
||||||
|
"expected /v1/chat/completions to be hit by the local-mode agent"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
+17
-1
@@ -13,6 +13,8 @@ use tower_http::services::ServeDir;
|
|||||||
mod admin;
|
mod admin;
|
||||||
mod auth;
|
mod auth;
|
||||||
mod checkin;
|
mod checkin;
|
||||||
|
#[cfg(feature = "multi_tenant_comfy")]
|
||||||
|
mod comfy_proxy;
|
||||||
mod comfyui;
|
mod comfyui;
|
||||||
mod comfyui_agent;
|
mod comfyui_agent;
|
||||||
mod events;
|
mod events;
|
||||||
@@ -28,6 +30,13 @@ mod tasks;
|
|||||||
mod update_center;
|
mod update_center;
|
||||||
mod user;
|
mod user;
|
||||||
mod user_api_keys;
|
mod user_api_keys;
|
||||||
|
// 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 video_agent;
|
||||||
|
|
||||||
pub fn build_router(state: AppState) -> Router {
|
pub fn build_router(state: AppState) -> Router {
|
||||||
@@ -39,7 +48,14 @@ pub fn build_router(state: AppState) -> Router {
|
|||||||
.merge(admin::routes(state.clone()))
|
.merge(admin::routes(state.clone()))
|
||||||
.merge(proxy::routes())
|
.merge(proxy::routes())
|
||||||
.merge(comfyui_agent::routes())
|
.merge(comfyui_agent::routes())
|
||||||
.merge(comfyui::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());
|
||||||
|
|
||||||
|
let api_routes = api_routes
|
||||||
.merge(auth::routes())
|
.merge(auth::routes())
|
||||||
.merge(events::routes())
|
.merge(events::routes())
|
||||||
.merge(checkin::routes())
|
.merge(checkin::routes())
|
||||||
|
|||||||
+309
-13
@@ -17,7 +17,7 @@ use crate::{
|
|||||||
AppState,
|
AppState,
|
||||||
};
|
};
|
||||||
use axum::{
|
use axum::{
|
||||||
extract::State,
|
extract::{Multipart, State},
|
||||||
http::HeaderMap,
|
http::HeaderMap,
|
||||||
response::{IntoResponse, Response, Sse},
|
response::{IntoResponse, Response, Sse},
|
||||||
routing::{get, post},
|
routing::{get, post},
|
||||||
@@ -25,7 +25,7 @@ use axum::{
|
|||||||
};
|
};
|
||||||
use futures::{stream, FutureExt, StreamExt};
|
use futures::{stream, FutureExt, StreamExt};
|
||||||
use serde_json::{json, Value};
|
use serde_json::{json, Value};
|
||||||
use std::{convert::Infallible, sync::Arc};
|
use std::{convert::Infallible, sync::Arc, time::Instant};
|
||||||
use tracing::info;
|
use tracing::info;
|
||||||
|
|
||||||
// ── Auth types ────────────────────────────────────────────────────────
|
// ── Auth types ────────────────────────────────────────────────────────
|
||||||
@@ -58,6 +58,7 @@ pub fn routes() -> Router<AppState> {
|
|||||||
.route("/v1/audio/translations", post(audio_translations))
|
.route("/v1/audio/translations", post(audio_translations))
|
||||||
.route("/v1/audio/speech", post(audio_speech))
|
.route("/v1/audio/speech", post(audio_speech))
|
||||||
.route("/v1/images/generations", post(images_generations))
|
.route("/v1/images/generations", post(images_generations))
|
||||||
|
.route("/v1/images/edits", post(images_edits))
|
||||||
.route("/v1/videos", post(videos_handler))
|
.route("/v1/videos", post(videos_handler))
|
||||||
.route("/v1/models", get(list_models))
|
.route("/v1/models", get(list_models))
|
||||||
// ── ComfyUI shim aliases ──────────────────────────────────────
|
// ── ComfyUI shim aliases ──────────────────────────────────────
|
||||||
@@ -70,6 +71,10 @@ pub fn routes() -> Router<AppState> {
|
|||||||
"/api/comfyui-agent/proxy/openai/images/generations",
|
"/api/comfyui-agent/proxy/openai/images/generations",
|
||||||
post(images_generations),
|
post(images_generations),
|
||||||
)
|
)
|
||||||
|
.route(
|
||||||
|
"/api/comfyui-agent/proxy/openai/images/edits",
|
||||||
|
post(images_edits),
|
||||||
|
)
|
||||||
.route(
|
.route(
|
||||||
"/api/comfyui-agent/proxy/openai/v1/responses",
|
"/api/comfyui-agent/proxy/openai/v1/responses",
|
||||||
post(responses_handler),
|
post(responses_handler),
|
||||||
@@ -321,6 +326,98 @@ async fn images_generations(
|
|||||||
Ok(Json(result))
|
Ok(Json(result))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// POST /v1/images/edits
|
||||||
|
///
|
||||||
|
/// `multipart/form-data` request: one or more image binaries
|
||||||
|
/// (`image`, `image[]`, `mask`) plus text fields (`prompt`, `model`,
|
||||||
|
/// `n`, `size`, `response_format`, `user`, `quality`, …). Mirrors
|
||||||
|
/// /images/generations for auth + routing + billing, but rebuilds a
|
||||||
|
/// fresh upstream multipart form rather than streaming the raw body —
|
||||||
|
/// that lets the route override `model` with the conductor-resolved
|
||||||
|
/// upstream model and skip frontend-only fields.
|
||||||
|
async fn images_edits(
|
||||||
|
State(state): State<AppState>,
|
||||||
|
headers: HeaderMap,
|
||||||
|
mut multipart: Multipart,
|
||||||
|
) -> AppResult<Json<Value>> {
|
||||||
|
info!("POST /v1/images/edits");
|
||||||
|
|
||||||
|
let key_id = check_auth(&state, &headers).await?;
|
||||||
|
|
||||||
|
let mut fields: Vec<executors::MultipartField> = Vec::new();
|
||||||
|
let mut model: Option<String> = None;
|
||||||
|
|
||||||
|
while let Some(field) = multipart
|
||||||
|
.next_field()
|
||||||
|
.await
|
||||||
|
.map_err(|e| AppError::BadRequest(format!("invalid multipart body: {e}")))?
|
||||||
|
{
|
||||||
|
let name = field
|
||||||
|
.name()
|
||||||
|
.ok_or_else(|| AppError::BadRequest("multipart field missing name".into()))?
|
||||||
|
.to_string();
|
||||||
|
let filename = field.file_name().map(str::to_string);
|
||||||
|
let content_type = field.content_type().map(str::to_string);
|
||||||
|
|
||||||
|
// OpenAI's spec lists exactly which fields carry binaries; everything
|
||||||
|
// else is a text field. We tag a field as binary when it has a
|
||||||
|
// filename _or_ a non-text content-type — covers both well-formed
|
||||||
|
// SDKs and minimal curl uploads.
|
||||||
|
let is_file = filename.is_some()
|
||||||
|
|| content_type
|
||||||
|
.as_deref()
|
||||||
|
.map(|ct| !ct.starts_with("text/") && ct != "application/x-www-form-urlencoded")
|
||||||
|
.unwrap_or(false);
|
||||||
|
|
||||||
|
if is_file {
|
||||||
|
let bytes = field
|
||||||
|
.bytes()
|
||||||
|
.await
|
||||||
|
.map_err(|e| AppError::BadRequest(format!("failed to read upload: {e}")))?
|
||||||
|
.to_vec();
|
||||||
|
fields.push(executors::MultipartField {
|
||||||
|
name,
|
||||||
|
bytes: Some(bytes),
|
||||||
|
filename,
|
||||||
|
content_type,
|
||||||
|
text: None,
|
||||||
|
});
|
||||||
|
} else {
|
||||||
|
let text = field
|
||||||
|
.text()
|
||||||
|
.await
|
||||||
|
.map_err(|e| AppError::BadRequest(format!("failed to read field: {e}")))?;
|
||||||
|
if name == "model" {
|
||||||
|
model = Some(text.clone());
|
||||||
|
}
|
||||||
|
fields.push(executors::MultipartField {
|
||||||
|
name,
|
||||||
|
bytes: None,
|
||||||
|
filename: None,
|
||||||
|
content_type: None,
|
||||||
|
text: Some(text),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
let model = model.ok_or_else(|| AppError::BadRequest("model is required".into()))?;
|
||||||
|
let route = resolve_route(&state, &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 {
|
||||||
|
for field in fields.iter_mut() {
|
||||||
|
if field.name == "model" && field.bytes.is_none() {
|
||||||
|
field.text = Some(route.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))
|
||||||
|
}
|
||||||
|
|
||||||
/// GET /v1/models
|
/// GET /v1/models
|
||||||
async fn list_models(State(state): State<AppState>, headers: HeaderMap) -> AppResult<Json<Value>> {
|
async fn list_models(State(state): State<AppState>, headers: HeaderMap) -> AppResult<Json<Value>> {
|
||||||
info!("GET /v1/models");
|
info!("GET /v1/models");
|
||||||
@@ -401,6 +498,91 @@ async fn record_billing(
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
async fn record_proxy_log_success(
|
||||||
|
state: &AppState,
|
||||||
|
key_id: Option<i64>,
|
||||||
|
selected: &SelectedChannel,
|
||||||
|
requested_model: &str,
|
||||||
|
latency_ms: Option<i64>,
|
||||||
|
prompt_tokens: i64,
|
||||||
|
completion_tokens: i64,
|
||||||
|
cost: f64,
|
||||||
|
attempts: usize,
|
||||||
|
is_stream: bool,
|
||||||
|
) {
|
||||||
|
if let Err(err) = sqlx::query(
|
||||||
|
"INSERT INTO proxy_logs (
|
||||||
|
route_id, channel_id, account_id, downstream_api_key_id,
|
||||||
|
model_requested, model_actual, status, http_status,
|
||||||
|
latency_ms, prompt_tokens, completion_tokens, total_tokens, estimated_cost,
|
||||||
|
retry_count, is_stream
|
||||||
|
) VALUES (?, ?, ?, ?, ?, ?, 'success', 200, ?, ?, ?, ?, ?, ?, ?)",
|
||||||
|
)
|
||||||
|
.bind(selected.route_id)
|
||||||
|
.bind(selected.channel_id)
|
||||||
|
.bind(selected.account_id)
|
||||||
|
.bind(key_id)
|
||||||
|
.bind(requested_model)
|
||||||
|
.bind(&selected.actual_model)
|
||||||
|
.bind(latency_ms)
|
||||||
|
.bind(prompt_tokens)
|
||||||
|
.bind(completion_tokens)
|
||||||
|
.bind(prompt_tokens + completion_tokens)
|
||||||
|
.bind(cost)
|
||||||
|
.bind(attempts.saturating_sub(1) as i64)
|
||||||
|
.bind(if is_stream { 1 } else { 0 })
|
||||||
|
.execute(&*state.db)
|
||||||
|
.await
|
||||||
|
{
|
||||||
|
tracing::warn!(
|
||||||
|
%err,
|
||||||
|
model = requested_model,
|
||||||
|
channel_id = selected.channel_id,
|
||||||
|
"failed to record proxy_logs success entry",
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn record_proxy_log_failure(
|
||||||
|
state: &AppState,
|
||||||
|
key_id: Option<i64>,
|
||||||
|
selected: Option<&SelectedChannel>,
|
||||||
|
requested_model: &str,
|
||||||
|
status: Option<u16>,
|
||||||
|
latency_ms: Option<i64>,
|
||||||
|
error_message: Option<&str>,
|
||||||
|
attempts: usize,
|
||||||
|
is_stream: bool,
|
||||||
|
) {
|
||||||
|
if let Err(err) = sqlx::query(
|
||||||
|
"INSERT INTO proxy_logs (
|
||||||
|
route_id, channel_id, account_id, downstream_api_key_id,
|
||||||
|
model_requested, model_actual, status, http_status,
|
||||||
|
latency_ms, error_message, retry_count, is_stream
|
||||||
|
) VALUES (?, ?, ?, ?, ?, ?, 'error', ?, ?, ?, ?, ?)",
|
||||||
|
)
|
||||||
|
.bind(selected.and_then(|item| item.route_id))
|
||||||
|
.bind(selected.map(|item| item.channel_id))
|
||||||
|
.bind(selected.and_then(|item| item.account_id))
|
||||||
|
.bind(key_id)
|
||||||
|
.bind(requested_model)
|
||||||
|
.bind(selected.map(|item| item.actual_model.as_str()))
|
||||||
|
.bind(status.map(i64::from))
|
||||||
|
.bind(latency_ms)
|
||||||
|
.bind(error_message.unwrap_or("proxy request failed"))
|
||||||
|
.bind(attempts.saturating_sub(1) as i64)
|
||||||
|
.bind(if is_stream { 1 } else { 0 })
|
||||||
|
.execute(&*state.db)
|
||||||
|
.await
|
||||||
|
{
|
||||||
|
tracing::warn!(
|
||||||
|
%err,
|
||||||
|
model = requested_model,
|
||||||
|
"failed to record proxy_logs failure entry",
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
/// Pull `usage` block out of an OpenAI-shaped response. Returns
|
/// Pull `usage` block out of an OpenAI-shaped response. Returns
|
||||||
/// (prompt_tokens, completion_tokens). Missing fields default to 0.
|
/// (prompt_tokens, completion_tokens). Missing fields default to 0.
|
||||||
fn extract_usage(response: &Value) -> (i64, i64) {
|
fn extract_usage(response: &Value) -> (i64, i64) {
|
||||||
@@ -419,6 +601,10 @@ fn extract_usage(response: &Value) -> (i64, i64) {
|
|||||||
(prompt, completion)
|
(prompt, completion)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn elapsed_ms_since(started_at: Instant) -> i64 {
|
||||||
|
started_at.elapsed().as_millis().min(i64::MAX as u128) as i64
|
||||||
|
}
|
||||||
|
|
||||||
// ── Route Resolution ──────────────────────────────────────────────────
|
// ── Route Resolution ──────────────────────────────────────────────────
|
||||||
|
|
||||||
/// Resolve upstream endpoint + API key for a given model.
|
/// Resolve upstream endpoint + API key for a given model.
|
||||||
@@ -491,6 +677,7 @@ async fn execute_chat_completions_with_conductor(
|
|||||||
model: String,
|
model: String,
|
||||||
body: Value,
|
body: Value,
|
||||||
) -> AppResult<Value> {
|
) -> AppResult<Value> {
|
||||||
|
let started_at = Instant::now();
|
||||||
let conductor = build_conductor(&state);
|
let conductor = build_conductor(&state);
|
||||||
let billing_state = state.clone();
|
let billing_state = state.clone();
|
||||||
let billing_model = model.clone();
|
let billing_model = model.clone();
|
||||||
@@ -527,10 +714,55 @@ async fn execute_chat_completions_with_conductor(
|
|||||||
})
|
})
|
||||||
.await;
|
.await;
|
||||||
|
|
||||||
let response = conductor_result_to_app_result(result)?;
|
match result {
|
||||||
let (prompt, completion) = extract_usage(&response);
|
ExecuteResult::Success {
|
||||||
record_billing(&billing_state, key_id, &billing_model, prompt, completion, 0.0).await;
|
selected,
|
||||||
Ok(response)
|
response,
|
||||||
|
attempts,
|
||||||
|
} => {
|
||||||
|
let (prompt, completion) = extract_usage(&response);
|
||||||
|
record_proxy_log_success(
|
||||||
|
&billing_state,
|
||||||
|
key_id,
|
||||||
|
&selected,
|
||||||
|
&billing_model,
|
||||||
|
Some(elapsed_ms_since(started_at)),
|
||||||
|
prompt,
|
||||||
|
completion,
|
||||||
|
0.0,
|
||||||
|
attempts,
|
||||||
|
false,
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
record_billing(&billing_state, key_id, &billing_model, prompt, completion, 0.0).await;
|
||||||
|
Ok(response)
|
||||||
|
}
|
||||||
|
ExecuteResult::Failure {
|
||||||
|
reason,
|
||||||
|
selected,
|
||||||
|
status,
|
||||||
|
raw_error_text,
|
||||||
|
attempts,
|
||||||
|
} => {
|
||||||
|
record_proxy_log_failure(
|
||||||
|
&billing_state,
|
||||||
|
key_id,
|
||||||
|
selected.as_ref(),
|
||||||
|
&billing_model,
|
||||||
|
status,
|
||||||
|
Some(elapsed_ms_since(started_at)),
|
||||||
|
raw_error_text.as_deref(),
|
||||||
|
attempts,
|
||||||
|
false,
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
Err(AppError::UpstreamError(format!(
|
||||||
|
"Proxy conductor failed after {attempts} attempts ({reason:?}, status: {:?}): {}",
|
||||||
|
status,
|
||||||
|
raw_error_text.unwrap_or_else(|| "no upstream error text".to_string())
|
||||||
|
)))
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn execute_chat_completions_stream_with_conductor(
|
async fn execute_chat_completions_stream_with_conductor(
|
||||||
@@ -539,6 +771,7 @@ async fn execute_chat_completions_stream_with_conductor(
|
|||||||
model: String,
|
model: String,
|
||||||
body: Value,
|
body: Value,
|
||||||
) -> AppResult<Vec<String>> {
|
) -> AppResult<Vec<String>> {
|
||||||
|
let started_at = Instant::now();
|
||||||
let conductor = build_conductor(&state);
|
let conductor = build_conductor(&state);
|
||||||
let billing_state = state.clone();
|
let billing_state = state.clone();
|
||||||
let billing_model = model.clone();
|
let billing_model = model.clone();
|
||||||
@@ -575,19 +808,50 @@ async fn execute_chat_completions_stream_with_conductor(
|
|||||||
})
|
})
|
||||||
.await;
|
.await;
|
||||||
|
|
||||||
let chunks = match conductor_result_to_app_result(result)? {
|
let (selected, attempts, chunks) = match result {
|
||||||
Value::Array(chunks) => chunks
|
ExecuteResult::Success {
|
||||||
|
selected,
|
||||||
|
response,
|
||||||
|
attempts,
|
||||||
|
} => match response {
|
||||||
|
Value::Array(chunks) => (selected, attempts, chunks
|
||||||
.into_iter()
|
.into_iter()
|
||||||
.map(|chunk| {
|
.map(|chunk| {
|
||||||
chunk.as_str().map(str::to_string).ok_or_else(|| {
|
chunk.as_str().map(str::to_string).ok_or_else(|| {
|
||||||
AppError::UpstreamError("Invalid streaming response from conductor".into())
|
AppError::UpstreamError("Invalid streaming response from conductor".into())
|
||||||
})
|
})
|
||||||
})
|
})
|
||||||
.collect::<AppResult<Vec<String>>>()?,
|
.collect::<AppResult<Vec<String>>>()?),
|
||||||
_ => {
|
_ => {
|
||||||
return Err(AppError::UpstreamError(
|
return Err(AppError::UpstreamError(
|
||||||
"Invalid streaming response from conductor".into(),
|
"Invalid streaming response from conductor".into(),
|
||||||
))
|
))
|
||||||
|
}
|
||||||
|
},
|
||||||
|
ExecuteResult::Failure {
|
||||||
|
reason,
|
||||||
|
selected,
|
||||||
|
status,
|
||||||
|
raw_error_text,
|
||||||
|
attempts,
|
||||||
|
} => {
|
||||||
|
record_proxy_log_failure(
|
||||||
|
&billing_state,
|
||||||
|
key_id,
|
||||||
|
selected.as_ref(),
|
||||||
|
&billing_model,
|
||||||
|
status,
|
||||||
|
Some(elapsed_ms_since(started_at)),
|
||||||
|
raw_error_text.as_deref(),
|
||||||
|
attempts,
|
||||||
|
true,
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
return Err(AppError::UpstreamError(format!(
|
||||||
|
"Proxy conductor failed after {attempts} attempts ({reason:?}, status: {:?}): {}",
|
||||||
|
status,
|
||||||
|
raw_error_text.unwrap_or_else(|| "no upstream error text".to_string())
|
||||||
|
)));
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
|
|
||||||
@@ -611,6 +875,19 @@ async fn execute_chat_completions_stream_with_conductor(
|
|||||||
}
|
}
|
||||||
})
|
})
|
||||||
.unwrap_or((0, 0));
|
.unwrap_or((0, 0));
|
||||||
|
record_proxy_log_success(
|
||||||
|
&billing_state,
|
||||||
|
key_id,
|
||||||
|
&selected,
|
||||||
|
&billing_model,
|
||||||
|
Some(elapsed_ms_since(started_at)),
|
||||||
|
prompt,
|
||||||
|
completion,
|
||||||
|
0.0,
|
||||||
|
attempts,
|
||||||
|
true,
|
||||||
|
)
|
||||||
|
.await;
|
||||||
record_billing(&billing_state, key_id, &billing_model, prompt, completion, 0.0).await;
|
record_billing(&billing_state, key_id, &billing_model, prompt, completion, 0.0).await;
|
||||||
Ok(chunks)
|
Ok(chunks)
|
||||||
}
|
}
|
||||||
@@ -881,6 +1158,25 @@ mod tests {
|
|||||||
.execute(&pool)
|
.execute(&pool)
|
||||||
.await
|
.await
|
||||||
.unwrap();
|
.unwrap();
|
||||||
|
// The channel-selection query joins through site_api_endpoints to
|
||||||
|
// pick a per-site override URL. We don't need any rows — the
|
||||||
|
// COALESCE falls back to sites.url when this table is empty —
|
||||||
|
// but the table itself must exist or the SELECT errors out and
|
||||||
|
// the conductor sees `None` (NoChannel) instead of an attempt
|
||||||
|
// failure that bumps fail_count.
|
||||||
|
sqlx::query(
|
||||||
|
"CREATE TABLE site_api_endpoints (
|
||||||
|
id INTEGER PRIMARY KEY,
|
||||||
|
site_id INTEGER NOT NULL,
|
||||||
|
url TEXT NOT NULL,
|
||||||
|
enabled INTEGER DEFAULT 1,
|
||||||
|
sort_order INTEGER DEFAULT 0,
|
||||||
|
cooldown_until TEXT
|
||||||
|
)",
|
||||||
|
)
|
||||||
|
.execute(&pool)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
sqlx::query("INSERT INTO downstream_api_keys (id, name, key, enabled) VALUES (1, 'test', 'sk-test', 1)")
|
sqlx::query("INSERT INTO downstream_api_keys (id, name, key, enabled) VALUES (1, 'test', 'sk-test', 1)")
|
||||||
.execute(&pool)
|
.execute(&pool)
|
||||||
|
|||||||
+52
-2
@@ -121,7 +121,29 @@ async fn proxy_logs(
|
|||||||
.max(0);
|
.max(0);
|
||||||
|
|
||||||
let items: Vec<Value> = sqlx::query_as::<_, ProxyLogRow>(
|
let items: Vec<Value> = sqlx::query_as::<_, ProxyLogRow>(
|
||||||
"SELECT id, route_id, channel_id, account_id, downstream_api_key_id, model_requested, model_actual, status, http_status, latency_ms, first_byte_latency_ms, prompt_tokens, completion_tokens, total_tokens, estimated_cost, error_message, retry_count, is_stream, client_family, client_app_id, client_app_name, client_confidence, created_at FROM proxy_logs ORDER BY created_at DESC LIMIT ? OFFSET ?",
|
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,
|
||||||
|
dk.name AS downstream_key_name,
|
||||||
|
dk.group_name AS downstream_key_group_name,
|
||||||
|
dk.tags AS downstream_key_tags
|
||||||
|
FROM proxy_logs l
|
||||||
|
LEFT JOIN accounts a ON a.id = l.account_id
|
||||||
|
LEFT JOIN sites s ON s.id = a.site_id
|
||||||
|
LEFT JOIN downstream_api_keys dk ON dk.id = l.downstream_api_key_id
|
||||||
|
ORDER BY l.created_at DESC
|
||||||
|
LIMIT ? OFFSET ?
|
||||||
|
"#,
|
||||||
)
|
)
|
||||||
.bind(limit)
|
.bind(limit)
|
||||||
.bind(offset)
|
.bind(offset)
|
||||||
@@ -162,6 +184,13 @@ pub struct ProxyLogRow {
|
|||||||
channel_id: Option<i64>,
|
channel_id: Option<i64>,
|
||||||
account_id: Option<i64>,
|
account_id: Option<i64>,
|
||||||
downstream_api_key_id: Option<i64>,
|
downstream_api_key_id: Option<i64>,
|
||||||
|
site_id: Option<i64>,
|
||||||
|
username: Option<String>,
|
||||||
|
site_name: Option<String>,
|
||||||
|
site_url: Option<String>,
|
||||||
|
downstream_key_name: Option<String>,
|
||||||
|
downstream_key_group_name: Option<String>,
|
||||||
|
downstream_key_tags: Option<String>,
|
||||||
model_requested: Option<String>,
|
model_requested: Option<String>,
|
||||||
model_actual: Option<String>,
|
model_actual: Option<String>,
|
||||||
status: Option<String>,
|
status: Option<String>,
|
||||||
@@ -187,7 +216,28 @@ async fn proxy_log_detail(
|
|||||||
axum::extract::Path(id): axum::extract::Path<i64>,
|
axum::extract::Path(id): axum::extract::Path<i64>,
|
||||||
) -> AppResult<Json<Value>> {
|
) -> AppResult<Json<Value>> {
|
||||||
let item: Option<ProxyLogRow> = sqlx::query_as(
|
let item: Option<ProxyLogRow> = sqlx::query_as(
|
||||||
"SELECT id, route_id, channel_id, account_id, downstream_api_key_id, model_requested, model_actual, status, http_status, latency_ms, first_byte_latency_ms, prompt_tokens, completion_tokens, total_tokens, estimated_cost, error_message, retry_count, is_stream, client_family, client_app_id, client_app_name, client_confidence, created_at FROM proxy_logs WHERE id = ?",
|
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,
|
||||||
|
dk.name AS downstream_key_name,
|
||||||
|
dk.group_name AS downstream_key_group_name,
|
||||||
|
dk.tags AS downstream_key_tags
|
||||||
|
FROM proxy_logs l
|
||||||
|
LEFT JOIN accounts a ON a.id = l.account_id
|
||||||
|
LEFT JOIN sites s ON s.id = a.site_id
|
||||||
|
LEFT JOIN downstream_api_keys dk ON dk.id = l.downstream_api_key_id
|
||||||
|
WHERE l.id = ?
|
||||||
|
"#,
|
||||||
)
|
)
|
||||||
.bind(id)
|
.bind(id)
|
||||||
.fetch_optional(&*state.db)
|
.fetch_optional(&*state.db)
|
||||||
|
|||||||
+204
-50
@@ -1,3 +1,15 @@
|
|||||||
|
//! `/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::{
|
use axum::{
|
||||||
body::Body,
|
body::Body,
|
||||||
extract::State,
|
extract::State,
|
||||||
@@ -12,13 +24,20 @@ use axum::{
|
|||||||
use futures::stream;
|
use futures::stream;
|
||||||
use serde_json::{json, Value};
|
use serde_json::{json, Value};
|
||||||
use std::convert::Infallible;
|
use std::convert::Infallible;
|
||||||
|
use std::time::{Duration, Instant};
|
||||||
use tracing::info;
|
use tracing::info;
|
||||||
|
|
||||||
use crate::{
|
use crate::{
|
||||||
|
auth::jwt::validate_jwt,
|
||||||
error::{AppError, AppResult},
|
error::{AppError, AppResult},
|
||||||
|
services::comfyui_pool::{ContainerState, PoolError},
|
||||||
AppState,
|
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 格式返回：
|
const SCRIPT_ANALYSIS_PROMPT: &str = r#"你是一个 AI 视频生成助手。分析用户提供的脚本，提取以下结构化信息，以 JSON 格式返回：
|
||||||
|
|
||||||
{
|
{
|
||||||
@@ -91,40 +110,156 @@ pub fn routes() -> Router<AppState> {
|
|||||||
.route("/api/video-agent/full-workflow", post(full_workflow))
|
.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>> {
|
async fn status(State(state): State<AppState>, headers: HeaderMap) -> AppResult<Json<Value>> {
|
||||||
check_auth(&headers, &state)?;
|
let user_id = resolve_user_id(&headers, &state)?;
|
||||||
Ok(Json(
|
Ok(Json(pool_status_json(&state, user_id).await))
|
||||||
serde_json::to_value(state.comfyui.status_response().await)
|
|
||||||
.unwrap_or_else(|_| json!({ "status": "unknown" })),
|
|
||||||
))
|
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn start_comfyui(
|
async fn start_comfyui(
|
||||||
State(state): State<AppState>,
|
State(state): State<AppState>,
|
||||||
headers: HeaderMap,
|
headers: HeaderMap,
|
||||||
) -> AppResult<Json<Value>> {
|
) -> AppResult<Json<Value>> {
|
||||||
check_auth(&headers, &state)?;
|
let user_id = resolve_user_id(&headers, &state)?;
|
||||||
|
// Idempotent: get_or_start returns the existing entry if already up.
|
||||||
state
|
state
|
||||||
.comfyui
|
.comfyui_pool
|
||||||
.start()
|
.get_or_start(user_id)
|
||||||
.await
|
.await
|
||||||
.map_err(|err| AppError::Internal(anyhow::anyhow!(err)))?;
|
.map_err(pool_err_to_app)?;
|
||||||
Ok(Json(json!({
|
Ok(Json(json!({
|
||||||
"success": true,
|
"success": true,
|
||||||
"comfyui": state.comfyui.status_response().await
|
"comfyui": pool_status_json(&state, user_id).await,
|
||||||
})))
|
})))
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn stop_comfyui(State(state): State<AppState>, headers: HeaderMap) -> AppResult<Json<Value>> {
|
async fn stop_comfyui(State(state): State<AppState>, headers: HeaderMap) -> AppResult<Json<Value>> {
|
||||||
check_auth(&headers, &state)?;
|
let user_id = resolve_user_id(&headers, &state)?;
|
||||||
state
|
state
|
||||||
.comfyui
|
.comfyui_pool
|
||||||
.stop()
|
.stop(user_id)
|
||||||
.await
|
.await
|
||||||
.map_err(|err| AppError::Internal(anyhow::anyhow!(err)))?;
|
.map_err(pool_err_to_app)?;
|
||||||
Ok(Json(json!({
|
Ok(Json(json!({
|
||||||
"success": true,
|
"success": true,
|
||||||
"comfyui": state.comfyui.status_response().await
|
"comfyui": pool_status_json(&state, user_id).await,
|
||||||
})))
|
})))
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -132,24 +267,45 @@ async fn restart_comfyui(
|
|||||||
State(state): State<AppState>,
|
State(state): State<AppState>,
|
||||||
headers: HeaderMap,
|
headers: HeaderMap,
|
||||||
) -> AppResult<Json<Value>> {
|
) -> AppResult<Json<Value>> {
|
||||||
check_auth(&headers, &state)?;
|
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
|
state
|
||||||
.comfyui
|
.comfyui_pool
|
||||||
.restart()
|
.get_or_start(user_id)
|
||||||
.await
|
.await
|
||||||
.map_err(|err| AppError::Internal(anyhow::anyhow!(err)))?;
|
.map_err(pool_err_to_app)?;
|
||||||
Ok(Json(json!({
|
Ok(Json(json!({
|
||||||
"success": true,
|
"success": true,
|
||||||
"comfyui": state.comfyui.status_response().await
|
"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(
|
async fn chat(
|
||||||
State(state): State<AppState>,
|
State(state): State<AppState>,
|
||||||
headers: HeaderMap,
|
headers: HeaderMap,
|
||||||
Json(body): Json<Value>,
|
Json(body): Json<Value>,
|
||||||
) -> AppResult<Sse<impl futures::Stream<Item = Result<Event, Infallible>>>> {
|
) -> AppResult<Sse<impl futures::Stream<Item = Result<Event, Infallible>>>> {
|
||||||
check_auth(&headers, &state)?;
|
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
|
let prompt = body
|
||||||
.get("message")
|
.get("message")
|
||||||
@@ -183,25 +339,6 @@ async fn chat(
|
|||||||
}))))
|
}))))
|
||||||
}
|
}
|
||||||
|
|
||||||
fn local_api_base(state: &AppState) -> String {
|
|
||||||
state.comfyui.model_backend_url()
|
|
||||||
}
|
|
||||||
|
|
||||||
fn check_auth(headers: &HeaderMap, state: &AppState) -> AppResult<()> {
|
|
||||||
let expected = format!("Bearer {}", state.config.auth.auth_token);
|
|
||||||
|
|
||||||
let actual = headers
|
|
||||||
.get(header::AUTHORIZATION)
|
|
||||||
.and_then(|v| v.to_str().ok())
|
|
||||||
.unwrap_or("");
|
|
||||||
|
|
||||||
if actual != expected {
|
|
||||||
return Err(AppError::BadRequest("unauthorized".to_string()));
|
|
||||||
}
|
|
||||||
|
|
||||||
Ok(())
|
|
||||||
}
|
|
||||||
|
|
||||||
fn build_basic_workflow(prompt: &str) -> Value {
|
fn build_basic_workflow(prompt: &str) -> Value {
|
||||||
json!({
|
json!({
|
||||||
"boosapi": {
|
"boosapi": {
|
||||||
@@ -255,7 +392,7 @@ fn parse_llm_json(raw: &str) -> AppResult<Value> {
|
|||||||
|
|
||||||
async fn call_llm(state: &AppState, prompt: &str, system: &str) -> AppResult<String> {
|
async fn call_llm(state: &AppState, prompt: &str, system: &str) -> AppResult<String> {
|
||||||
let client = reqwest::Client::new();
|
let client = reqwest::Client::new();
|
||||||
let url = format!("{}/chat/completions", local_api_base(state));
|
let url = format!("{}/chat/completions", boosapi_v1_base(state));
|
||||||
|
|
||||||
info!("video-agent calling BoosAPI LLM endpoint: {}", url);
|
info!("video-agent calling BoosAPI LLM endpoint: {}", url);
|
||||||
|
|
||||||
@@ -301,7 +438,7 @@ async fn call_llm(state: &AppState, prompt: &str, system: &str) -> AppResult<Str
|
|||||||
|
|
||||||
async fn call_image_gen(state: &AppState, prompt: &str) -> AppResult<String> {
|
async fn call_image_gen(state: &AppState, prompt: &str) -> AppResult<String> {
|
||||||
let client = reqwest::Client::new();
|
let client = reqwest::Client::new();
|
||||||
let url = format!("{}/images/generations", local_api_base(state));
|
let url = format!("{}/images/generations", boosapi_v1_base(state));
|
||||||
|
|
||||||
info!("video-agent calling BoosAPI image endpoint: {}", url);
|
info!("video-agent calling BoosAPI image endpoint: {}", url);
|
||||||
|
|
||||||
@@ -344,7 +481,7 @@ async fn call_image_gen(state: &AppState, prompt: &str) -> AppResult<String> {
|
|||||||
|
|
||||||
async fn call_tts(state: &AppState, text: &str, voice: &str) -> AppResult<Vec<u8>> {
|
async fn call_tts(state: &AppState, text: &str, voice: &str) -> AppResult<Vec<u8>> {
|
||||||
let client = reqwest::Client::new();
|
let client = reqwest::Client::new();
|
||||||
let url = format!("{}/audio/speech", local_api_base(state));
|
let url = format!("{}/audio/speech", boosapi_v1_base(state));
|
||||||
|
|
||||||
info!("video-agent calling BoosAPI TTS endpoint: {}", url);
|
info!("video-agent calling BoosAPI TTS endpoint: {}", url);
|
||||||
|
|
||||||
@@ -379,12 +516,15 @@ async fn call_tts(state: &AppState, text: &str, voice: &str) -> AppResult<Vec<u8
|
|||||||
Ok(bytes.to_vec())
|
Ok(bytes.to_vec())
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// ── Per-request handlers (5.5) ─────────────────────────────────────────
|
||||||
|
|
||||||
async fn analyze(
|
async fn analyze(
|
||||||
State(state): State<AppState>,
|
State(state): State<AppState>,
|
||||||
headers: HeaderMap,
|
headers: HeaderMap,
|
||||||
Json(body): Json<Value>,
|
Json(body): Json<Value>,
|
||||||
) -> AppResult<Json<Value>> {
|
) -> AppResult<Json<Value>> {
|
||||||
check_auth(&headers, &state)?;
|
let user_id = resolve_user_id(&headers, &state)?;
|
||||||
|
state.comfyui_pool.touch(user_id).await;
|
||||||
|
|
||||||
let script = required_str(&body, "script")?;
|
let script = required_str(&body, "script")?;
|
||||||
|
|
||||||
@@ -402,7 +542,8 @@ async fn character_prompts(
|
|||||||
headers: HeaderMap,
|
headers: HeaderMap,
|
||||||
Json(body): Json<Value>,
|
Json(body): Json<Value>,
|
||||||
) -> AppResult<Json<Value>> {
|
) -> AppResult<Json<Value>> {
|
||||||
check_auth(&headers, &state)?;
|
let user_id = resolve_user_id(&headers, &state)?;
|
||||||
|
state.comfyui_pool.touch(user_id).await;
|
||||||
|
|
||||||
let character = required_value(&body, "character")?;
|
let character = required_value(&body, "character")?;
|
||||||
|
|
||||||
@@ -438,7 +579,10 @@ async fn generate_character(
|
|||||||
headers: HeaderMap,
|
headers: HeaderMap,
|
||||||
Json(body): Json<Value>,
|
Json(body): Json<Value>,
|
||||||
) -> AppResult<Json<Value>> {
|
) -> AppResult<Json<Value>> {
|
||||||
check_auth(&headers, &state)?;
|
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 prompts = required_value(&body, "prompts")?;
|
||||||
let prompt_obj = prompts
|
let prompt_obj = prompts
|
||||||
@@ -456,6 +600,8 @@ async fn generate_character(
|
|||||||
results.insert(angle.clone(), Value::String(url));
|
results.insert(angle.clone(), Value::String(url));
|
||||||
}
|
}
|
||||||
|
|
||||||
|
state.comfyui_pool.touch(user_id).await;
|
||||||
|
|
||||||
Ok(Json(json!({
|
Ok(Json(json!({
|
||||||
"success": true,
|
"success": true,
|
||||||
"data": Value::Object(results)
|
"data": Value::Object(results)
|
||||||
@@ -467,7 +613,8 @@ async fn scene_prompt(
|
|||||||
headers: HeaderMap,
|
headers: HeaderMap,
|
||||||
Json(body): Json<Value>,
|
Json(body): Json<Value>,
|
||||||
) -> AppResult<Json<Value>> {
|
) -> AppResult<Json<Value>> {
|
||||||
check_auth(&headers, &state)?;
|
let user_id = resolve_user_id(&headers, &state)?;
|
||||||
|
state.comfyui_pool.touch(user_id).await;
|
||||||
|
|
||||||
let scene = required_value(&body, "scene")?;
|
let scene = required_value(&body, "scene")?;
|
||||||
|
|
||||||
@@ -495,11 +642,14 @@ async fn generate_scene(
|
|||||||
headers: HeaderMap,
|
headers: HeaderMap,
|
||||||
Json(body): Json<Value>,
|
Json(body): Json<Value>,
|
||||||
) -> AppResult<Json<Value>> {
|
) -> AppResult<Json<Value>> {
|
||||||
check_auth(&headers, &state)?;
|
let user_id = resolve_user_id(&headers, &state)?;
|
||||||
|
let _ = ensure_container_ready(&state, user_id).await?;
|
||||||
|
|
||||||
let prompt = required_str(&body, "prompt")?;
|
let prompt = required_str(&body, "prompt")?;
|
||||||
let url = call_image_gen(&state, prompt).await?;
|
let url = call_image_gen(&state, prompt).await?;
|
||||||
|
|
||||||
|
state.comfyui_pool.touch(user_id).await;
|
||||||
|
|
||||||
Ok(Json(json!({
|
Ok(Json(json!({
|
||||||
"success": true,
|
"success": true,
|
||||||
"data": {
|
"data": {
|
||||||
@@ -513,7 +663,8 @@ async fn generate_tts(
|
|||||||
headers: HeaderMap,
|
headers: HeaderMap,
|
||||||
Json(body): Json<Value>,
|
Json(body): Json<Value>,
|
||||||
) -> AppResult<Response<Body>> {
|
) -> AppResult<Response<Body>> {
|
||||||
check_auth(&headers, &state)?;
|
let user_id = resolve_user_id(&headers, &state)?;
|
||||||
|
state.comfyui_pool.touch(user_id).await;
|
||||||
|
|
||||||
let text = required_str(&body, "text")?;
|
let text = required_str(&body, "text")?;
|
||||||
let voice = body
|
let voice = body
|
||||||
@@ -537,7 +688,8 @@ async fn full_workflow(
|
|||||||
headers: HeaderMap,
|
headers: HeaderMap,
|
||||||
Json(body): Json<Value>,
|
Json(body): Json<Value>,
|
||||||
) -> AppResult<Json<Value>> {
|
) -> AppResult<Json<Value>> {
|
||||||
check_auth(&headers, &state)?;
|
let user_id = resolve_user_id(&headers, &state)?;
|
||||||
|
let _ = ensure_container_ready(&state, user_id).await?;
|
||||||
|
|
||||||
let script = required_str(&body, "script")?;
|
let script = required_str(&body, "script")?;
|
||||||
let requested_style = body.get("style").and_then(Value::as_str);
|
let requested_style = body.get("style").and_then(Value::as_str);
|
||||||
@@ -611,6 +763,8 @@ async fn full_workflow(
|
|||||||
.or_else(|| requested_style.map(|s| Value::String(s.to_string())))
|
.or_else(|| requested_style.map(|s| Value::String(s.to_string())))
|
||||||
.unwrap_or(Value::Null);
|
.unwrap_or(Value::Null);
|
||||||
|
|
||||||
|
state.comfyui_pool.touch(user_id).await;
|
||||||
|
|
||||||
Ok(Json(json!({
|
Ok(Json(json!({
|
||||||
"success": true,
|
"success": true,
|
||||||
"data": {
|
"data": {
|
||||||
|
|||||||
@@ -0,0 +1,623 @@
|
|||||||
|
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 tracing::info;
|
||||||
|
|
||||||
|
use crate::{
|
||||||
|
error::{AppError, AppResult},
|
||||||
|
AppState,
|
||||||
|
};
|
||||||
|
|
||||||
|
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))
|
||||||
|
}
|
||||||
|
|
||||||
|
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" })),
|
||||||
|
))
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn start_comfyui(
|
||||||
|
State(state): State<AppState>,
|
||||||
|
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
|
||||||
|
})))
|
||||||
|
}
|
||||||
|
|
||||||
|
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
|
||||||
|
})))
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn restart_comfyui(
|
||||||
|
State(state): State<AppState>,
|
||||||
|
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
|
||||||
|
})))
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn chat(
|
||||||
|
State(state): State<AppState>,
|
||||||
|
headers: HeaderMap,
|
||||||
|
Json(body): Json<Value>,
|
||||||
|
) -> AppResult<Sse<impl futures::Stream<Item = Result<Event, Infallible>>>> {
|
||||||
|
check_auth(&headers, &state)?;
|
||||||
|
|
||||||
|
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 local_api_base(state: &AppState) -> String {
|
||||||
|
state.comfyui.model_backend_url()
|
||||||
|
}
|
||||||
|
|
||||||
|
fn check_auth(headers: &HeaderMap, state: &AppState) -> AppResult<()> {
|
||||||
|
let expected = format!("Bearer {}", state.config.auth.auth_token);
|
||||||
|
|
||||||
|
let actual = headers
|
||||||
|
.get(header::AUTHORIZATION)
|
||||||
|
.and_then(|v| v.to_str().ok())
|
||||||
|
.unwrap_or("");
|
||||||
|
|
||||||
|
if actual != expected {
|
||||||
|
return Err(AppError::BadRequest("unauthorized".to_string()));
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
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", local_api_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", local_api_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", local_api_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())
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn analyze(
|
||||||
|
State(state): State<AppState>,
|
||||||
|
headers: HeaderMap,
|
||||||
|
Json(body): Json<Value>,
|
||||||
|
) -> AppResult<Json<Value>> {
|
||||||
|
check_auth(&headers, &state)?;
|
||||||
|
|
||||||
|
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>> {
|
||||||
|
check_auth(&headers, &state)?;
|
||||||
|
|
||||||
|
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>> {
|
||||||
|
check_auth(&headers, &state)?;
|
||||||
|
|
||||||
|
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));
|
||||||
|
}
|
||||||
|
|
||||||
|
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>> {
|
||||||
|
check_auth(&headers, &state)?;
|
||||||
|
|
||||||
|
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>> {
|
||||||
|
check_auth(&headers, &state)?;
|
||||||
|
|
||||||
|
let prompt = required_str(&body, "prompt")?;
|
||||||
|
let url = call_image_gen(&state, prompt).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>> {
|
||||||
|
check_auth(&headers, &state)?;
|
||||||
|
|
||||||
|
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>> {
|
||||||
|
check_auth(&headers, &state)?;
|
||||||
|
|
||||||
|
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);
|
||||||
|
|
||||||
|
Ok(Json(json!({
|
||||||
|
"success": true,
|
||||||
|
"data": {
|
||||||
|
"title": analysis.get("title").cloned().unwrap_or(Value::Null),
|
||||||
|
"style": style,
|
||||||
|
"characters": characters,
|
||||||
|
"scenes": scenes
|
||||||
|
}
|
||||||
|
})))
|
||||||
|
}
|
||||||
@@ -0,0 +1,654 @@
|
|||||||
|
//! 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.
|
||||||
|
if let Some(entry) = self.containers.read().await.get(&user_id).cloned() {
|
||||||
|
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.
|
||||||
|
if let Some(entry) = self.containers.read().await.get(&user_id).cloned() {
|
||||||
|
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(), self.config.boosapi_base.clone());
|
||||||
|
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(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 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);
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -0,0 +1,120 @@
|
|||||||
|
//! 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>>;
|
||||||
|
}
|
||||||
@@ -0,0 +1,288 @@
|
|||||||
|
//! 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
|
||||||
|
))
|
||||||
|
// 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,9 +1,13 @@
|
|||||||
pub mod announcement;
|
pub mod announcement;
|
||||||
pub mod billing;
|
pub mod billing;
|
||||||
pub mod builtin_keys;
|
pub mod builtin_keys;
|
||||||
|
pub mod comfyui_pool;
|
||||||
pub mod comfyui_process;
|
pub mod comfyui_process;
|
||||||
|
pub mod container_runtime;
|
||||||
pub mod cron;
|
pub mod cron;
|
||||||
|
pub mod docker_runtime;
|
||||||
pub mod notification;
|
pub mod notification;
|
||||||
pub mod oauth;
|
pub mod oauth;
|
||||||
|
pub mod route_rebuild;
|
||||||
pub mod shared;
|
pub mod shared;
|
||||||
pub mod update_center;
|
pub mod update_center;
|
||||||
|
|||||||
@@ -13,6 +13,14 @@
|
|||||||
//! marked `manual_override = 1` are NEVER touched (created, updated, deleted).
|
//! marked `manual_override = 1` are NEVER touched (created, updated, deleted).
|
||||||
//! Other channels are reconciled to match the current model availability:
|
//! Other channels are reconciled to match the current model availability:
|
||||||
//! missing combos get inserted, stale combos get hard-deleted.
|
//! missing combos get inserted, stale combos get hard-deleted.
|
||||||
|
//! - After channels are reconciled, a route orphan sweep deletes
|
||||||
|
//! `token_routes` rows that (a) no active account advertises any more,
|
||||||
|
//! (b) have no surviving channels (so no manual_override channel pins
|
||||||
|
//! them), and (c) have empty `display_name` / `display_icon` /
|
||||||
|
//! `model_mapping` (so no operator has touched the route's display
|
||||||
|
//! metadata). Without this sweep, `/api/routes/summary` keeps surfacing
|
||||||
|
//! phantom routes for models nothing can serve, and the user-visible
|
||||||
|
//! "rebuild" button appears to do nothing.
|
||||||
//! - Disabled site models (`site_disabled_models`) are filtered out before the
|
//! - Disabled site models (`site_disabled_models`) are filtered out before the
|
||||||
//! write. So is anything from accounts whose status isn't `active`, or models
|
//! write. So is anything from accounts whose status isn't `active`, or models
|
||||||
//! where `model_availability.available != 1`.
|
//! where `model_availability.available != 1`.
|
||||||
@@ -30,6 +38,7 @@ use std::collections::{HashMap, HashSet};
|
|||||||
pub struct RebuildSummary {
|
pub struct RebuildSummary {
|
||||||
pub created_routes: i64,
|
pub created_routes: i64,
|
||||||
pub reused_routes: i64,
|
pub reused_routes: i64,
|
||||||
|
pub deleted_routes: i64,
|
||||||
pub created_channels: i64,
|
pub created_channels: i64,
|
||||||
pub deleted_channels: i64,
|
pub deleted_channels: i64,
|
||||||
pub skipped_manual_channels: i64,
|
pub skipped_manual_channels: i64,
|
||||||
@@ -255,6 +264,50 @@ pub(crate) async fn rebuild_routes_inner(
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// 6. Route orphan sweep. After channels for "no longer routable" models
|
||||||
|
// have been deleted above, the parent `token_routes` row can also
|
||||||
|
// safely go away — but only if it's a pure auto-managed leftover.
|
||||||
|
// We preserve a route when ANY of these is true:
|
||||||
|
// - the model is still in `by_model` (an active account advertises it);
|
||||||
|
// - the route still has at least one channel left, which after the
|
||||||
|
// channel sweep above can only be a `manual_override = 1` channel;
|
||||||
|
// - an operator has touched the route's display metadata
|
||||||
|
// (`display_name`, `display_icon`, `model_mapping`).
|
||||||
|
// Without this, every route ever created sticks around forever — the
|
||||||
|
// `/api/routes/summary` endpoint then surfaces phantom routes for
|
||||||
|
// models nothing can serve, and rebuild "does nothing" because the
|
||||||
|
// main loop only visits models that are in `by_model`.
|
||||||
|
let route_orphan_candidates: Vec<(i64, Option<String>, Option<String>, Option<String>)> =
|
||||||
|
sqlx::query_as(
|
||||||
|
"SELECT tr.id, tr.display_name, tr.display_icon, tr.model_mapping
|
||||||
|
FROM token_routes tr
|
||||||
|
WHERE NOT EXISTS (
|
||||||
|
SELECT 1 FROM route_channels rc WHERE rc.route_id = tr.id
|
||||||
|
)",
|
||||||
|
)
|
||||||
|
.fetch_all(&mut *tx)
|
||||||
|
.await?;
|
||||||
|
|
||||||
|
for (route_id, display_name, display_icon, model_mapping) in route_orphan_candidates {
|
||||||
|
if processed_route_ids.contains(&route_id) {
|
||||||
|
// Still backed by an active model — leave it, even if it currently
|
||||||
|
// has no channels (channels will be re-inserted on next refresh).
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
let touched_by_operator = display_name.as_deref().map(str::trim).unwrap_or("").is_empty()
|
||||||
|
== false
|
||||||
|
|| display_icon.as_deref().map(str::trim).unwrap_or("").is_empty() == false
|
||||||
|
|| model_mapping.as_deref().map(str::trim).unwrap_or("").is_empty() == false;
|
||||||
|
if touched_by_operator {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
sqlx::query("DELETE FROM token_routes WHERE id = ?")
|
||||||
|
.bind(route_id)
|
||||||
|
.execute(&mut *tx)
|
||||||
|
.await?;
|
||||||
|
summary.deleted_routes += 1;
|
||||||
|
}
|
||||||
|
|
||||||
tx.commit().await?;
|
tx.commit().await?;
|
||||||
Ok(summary)
|
Ok(summary)
|
||||||
}
|
}
|
||||||
@@ -296,6 +349,9 @@ mod tests {
|
|||||||
"CREATE TABLE token_routes (
|
"CREATE TABLE token_routes (
|
||||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||||
model_pattern TEXT NOT NULL UNIQUE,
|
model_pattern TEXT NOT NULL UNIQUE,
|
||||||
|
model_mapping TEXT,
|
||||||
|
display_name TEXT,
|
||||||
|
display_icon TEXT,
|
||||||
routing_strategy TEXT NOT NULL DEFAULT 'weighted',
|
routing_strategy TEXT NOT NULL DEFAULT 'weighted',
|
||||||
route_mode TEXT NOT NULL DEFAULT 'pattern',
|
route_mode TEXT NOT NULL DEFAULT 'pattern',
|
||||||
enabled INTEGER NOT NULL DEFAULT 1
|
enabled INTEGER NOT NULL DEFAULT 1
|
||||||
@@ -387,14 +443,23 @@ mod tests {
|
|||||||
summary.deleted_channels, 1,
|
summary.deleted_channels, 1,
|
||||||
"summary should report exactly one deletion"
|
"summary should report exactly one deletion"
|
||||||
);
|
);
|
||||||
// The route row itself stays — operators may have edited display
|
// The route row also goes — it's a pure auto-managed leftover (no
|
||||||
// metadata on it, and the next time an active account advertises
|
// surviving channels, no operator-edited display metadata, no active
|
||||||
// `gpt-x` we want to reuse it. The rebuild never deletes routes.
|
// account advertising the model). Leaving it behind is what causes
|
||||||
|
// `/api/routes/summary` to surface phantom routes for models nothing
|
||||||
|
// can serve.
|
||||||
let route_count: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM token_routes")
|
let route_count: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM token_routes")
|
||||||
.fetch_one(&pool)
|
.fetch_one(&pool)
|
||||||
.await
|
.await
|
||||||
.unwrap();
|
.unwrap();
|
||||||
assert_eq!(route_count, 1, "token_routes row should be reused, not dropped");
|
assert_eq!(
|
||||||
|
route_count, 0,
|
||||||
|
"pure auto-managed route must be swept once its last channel is gone"
|
||||||
|
);
|
||||||
|
assert_eq!(
|
||||||
|
summary.deleted_routes, 1,
|
||||||
|
"summary should report the route deletion"
|
||||||
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Sanity check: manual_override channels survive an orphan sweep. Even
|
/// Sanity check: manual_override channels survive an orphan sweep. Even
|
||||||
@@ -423,6 +488,10 @@ mod tests {
|
|||||||
|
|
||||||
let summary = rebuild_routes_inner(&pool).await.expect("rebuild");
|
let summary = rebuild_routes_inner(&pool).await.expect("rebuild");
|
||||||
assert_eq!(summary.deleted_channels, 0, "manual channel must survive");
|
assert_eq!(summary.deleted_channels, 0, "manual channel must survive");
|
||||||
|
assert_eq!(
|
||||||
|
summary.deleted_routes, 0,
|
||||||
|
"route is pinned by a manual_override channel and must survive"
|
||||||
|
);
|
||||||
|
|
||||||
let manual_count: i64 = sqlx::query_scalar(
|
let manual_count: i64 = sqlx::query_scalar(
|
||||||
"SELECT COUNT(*) FROM route_channels WHERE manual_override = 1",
|
"SELECT COUNT(*) FROM route_channels WHERE manual_override = 1",
|
||||||
@@ -431,5 +500,106 @@ mod tests {
|
|||||||
.await
|
.await
|
||||||
.unwrap();
|
.unwrap();
|
||||||
assert_eq!(manual_count, 1);
|
assert_eq!(manual_count, 1);
|
||||||
|
let route_count: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM token_routes")
|
||||||
|
.fetch_one(&pool)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
assert_eq!(route_count, 1);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// User-reported bug: empty DB still shows hundreds of routes on the
|
||||||
|
/// `/routes` page, and clicking "rebuild" does nothing. Reproduces by
|
||||||
|
/// pre-seeding `token_routes` rows that no account / availability row
|
||||||
|
/// references — exactly the state we found in the live `boosapi.db`.
|
||||||
|
/// The rebuild must hard-delete those phantom routes.
|
||||||
|
#[tokio::test]
|
||||||
|
async fn phantom_routes_with_no_accounts_are_swept_on_rebuild() {
|
||||||
|
let pool = fixture_pool().await;
|
||||||
|
|
||||||
|
// No accounts, no availability — but three left-behind auto-managed
|
||||||
|
// routes. Mirror the production layout: empty display_name /
|
||||||
|
// display_icon / model_mapping.
|
||||||
|
sqlx::query(
|
||||||
|
"INSERT INTO token_routes (id, model_pattern) VALUES
|
||||||
|
(1, 'gpt-4'),
|
||||||
|
(2, 'claude-opus-4.5'),
|
||||||
|
(3, 'glm-5')",
|
||||||
|
)
|
||||||
|
.execute(&pool)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
let summary = rebuild_routes_inner(&pool).await.expect("rebuild");
|
||||||
|
|
||||||
|
let route_count: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM token_routes")
|
||||||
|
.fetch_one(&pool)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
assert_eq!(
|
||||||
|
route_count, 0,
|
||||||
|
"all phantom routes must be swept when no account advertises any model"
|
||||||
|
);
|
||||||
|
assert_eq!(summary.deleted_routes, 3, "summary reports the three sweeps");
|
||||||
|
assert_eq!(summary.created_routes, 0);
|
||||||
|
assert_eq!(summary.created_channels, 0);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Operator-edited routes survive the sweep. We probe each of the three
|
||||||
|
/// "touched" columns (`display_name`, `display_icon`, `model_mapping`)
|
||||||
|
/// independently so an accidental refactor that drops a check in the
|
||||||
|
/// disjunction gets caught.
|
||||||
|
#[tokio::test]
|
||||||
|
async fn route_with_operator_metadata_is_preserved() {
|
||||||
|
for (label, sql) in [
|
||||||
|
(
|
||||||
|
"display_name",
|
||||||
|
"INSERT INTO token_routes (id, model_pattern, display_name) VALUES (1, 'gpt-4', 'Renamed by ops')",
|
||||||
|
),
|
||||||
|
(
|
||||||
|
"display_icon",
|
||||||
|
"INSERT INTO token_routes (id, model_pattern, display_icon) VALUES (1, 'gpt-4', 'openai-color')",
|
||||||
|
),
|
||||||
|
(
|
||||||
|
"model_mapping",
|
||||||
|
"INSERT INTO token_routes (id, model_pattern, model_mapping) VALUES (1, 'gpt-4', '{\"alias\":\"gpt-4o\"}')",
|
||||||
|
),
|
||||||
|
] {
|
||||||
|
let pool = fixture_pool().await;
|
||||||
|
sqlx::query(sql).execute(&pool).await.unwrap();
|
||||||
|
|
||||||
|
let summary = rebuild_routes_inner(&pool).await.expect("rebuild");
|
||||||
|
assert_eq!(
|
||||||
|
summary.deleted_routes, 0,
|
||||||
|
"operator-edited {label} must pin the route through a sweep"
|
||||||
|
);
|
||||||
|
let route_count: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM token_routes")
|
||||||
|
.fetch_one(&pool)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
assert_eq!(route_count, 1, "route must survive ({label} edited)");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Whitespace-only metadata is not an "edit" — `' '` should still get
|
||||||
|
/// swept. Otherwise a stray space saved by the admin form would
|
||||||
|
/// silently keep dead routes around forever.
|
||||||
|
#[tokio::test]
|
||||||
|
async fn whitespace_only_metadata_does_not_pin_a_route() {
|
||||||
|
let pool = fixture_pool().await;
|
||||||
|
sqlx::query(
|
||||||
|
"INSERT INTO token_routes (id, model_pattern, display_name, display_icon, model_mapping)
|
||||||
|
VALUES (1, 'gpt-4', ' ', '', '\t')",
|
||||||
|
)
|
||||||
|
.execute(&pool)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
|
||||||
|
let summary = rebuild_routes_inner(&pool).await.expect("rebuild");
|
||||||
|
assert_eq!(summary.deleted_routes, 1);
|
||||||
|
let route_count: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM token_routes")
|
||||||
|
.fetch_one(&pool)
|
||||||
|
.await
|
||||||
|
.unwrap();
|
||||||
|
assert_eq!(route_count, 0);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -1,7 +1,7 @@
|
|||||||
import { describe, expect, it, vi } from 'vitest';
|
import { describe, expect, it, vi } from 'vitest';
|
||||||
import { create, type ReactTestInstance } from 'react-test-renderer';
|
import { create, type ReactTestInstance } from 'react-test-renderer';
|
||||||
import { MemoryRouter } from 'react-router-dom';
|
import { MemoryRouter } from 'react-router-dom';
|
||||||
import { Login } from './App.js';
|
import AdminLoginPage from './pages/auth/AdminLoginPage.js';
|
||||||
import { SITE_DOCS_URL, SITE_GITHUB_URL } from './docsLink.js';
|
import { SITE_DOCS_URL, SITE_GITHUB_URL } from './docsLink.js';
|
||||||
|
|
||||||
function collectText(node: ReactTestInstance): string {
|
function collectText(node: ReactTestInstance): string {
|
||||||
@@ -23,7 +23,7 @@ describe('Login surface', () => {
|
|||||||
it('renders a poster-style hero with a floating admin login panel', () => {
|
it('renders a poster-style hero with a floating admin login panel', () => {
|
||||||
const root = create(
|
const root = create(
|
||||||
<MemoryRouter>
|
<MemoryRouter>
|
||||||
<Login onLogin={vi.fn()} t={(text) => text} initialAudience="admin" />
|
<AdminLoginPage onLogin={vi.fn()} t={(text) => text} />
|
||||||
</MemoryRouter>,
|
</MemoryRouter>,
|
||||||
);
|
);
|
||||||
|
|
||||||
@@ -79,7 +79,7 @@ describe('Login surface', () => {
|
|||||||
expect(githubLink.props['aria-label']).toBe('GitHub');
|
expect(githubLink.props['aria-label']).toBe('GitHub');
|
||||||
expect(githubLink.props.target).toBe('_blank');
|
expect(githubLink.props.target).toBe('_blank');
|
||||||
expect(tokenInput.props.type).toBe('password');
|
expect(tokenInput.props.type).toBe('password');
|
||||||
expect(loginModeTabs.map((node) => collectText(node))).toEqual(['普通用户', '管理员', 'Auth Token', '账号密码']);
|
expect(loginModeTabs.map((node) => collectText(node))).toEqual(['Auth Token', '账号密码']);
|
||||||
} finally {
|
} finally {
|
||||||
root?.unmount();
|
root?.unmount();
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -12,6 +12,9 @@ const { apiMock, authSessionMock } = vi.hoisted(() => ({
|
|||||||
hasValidAuthSession: vi.fn(),
|
hasValidAuthSession: vi.fn(),
|
||||||
persistAuthSession: vi.fn(),
|
persistAuthSession: vi.fn(),
|
||||||
clearAuthSession: vi.fn(),
|
clearAuthSession: vi.fn(),
|
||||||
|
getSessionToken: vi.fn(),
|
||||||
|
commitSession: vi.fn(),
|
||||||
|
clearSession: vi.fn(),
|
||||||
},
|
},
|
||||||
}));
|
}));
|
||||||
|
|
||||||
@@ -31,6 +34,9 @@ vi.mock('./authSession.js', () => ({
|
|||||||
hasValidAuthSession: authSessionMock.hasValidAuthSession,
|
hasValidAuthSession: authSessionMock.hasValidAuthSession,
|
||||||
persistAuthSession: authSessionMock.persistAuthSession,
|
persistAuthSession: authSessionMock.persistAuthSession,
|
||||||
clearAuthSession: authSessionMock.clearAuthSession,
|
clearAuthSession: authSessionMock.clearAuthSession,
|
||||||
|
getSessionToken: authSessionMock.getSessionToken,
|
||||||
|
commitSession: authSessionMock.commitSession,
|
||||||
|
clearSession: authSessionMock.clearSession,
|
||||||
}));
|
}));
|
||||||
|
|
||||||
vi.mock('./components/SearchModal.js', () => ({
|
vi.mock('./components/SearchModal.js', () => ({
|
||||||
@@ -138,6 +144,9 @@ describe('App mobile layout', () => {
|
|||||||
vi.useFakeTimers();
|
vi.useFakeTimers();
|
||||||
apiMock.getEvents.mockResolvedValue([]);
|
apiMock.getEvents.mockResolvedValue([]);
|
||||||
authSessionMock.hasValidAuthSession.mockReturnValue(true);
|
authSessionMock.hasValidAuthSession.mockReturnValue(true);
|
||||||
|
authSessionMock.getSessionToken.mockImplementation((audience: string) =>
|
||||||
|
audience === 'admin' ? 'fake-admin-token' : null,
|
||||||
|
);
|
||||||
});
|
});
|
||||||
|
|
||||||
afterEach(() => {
|
afterEach(() => {
|
||||||
|
|||||||
@@ -12,6 +12,9 @@ const { apiMock, authSessionMock } = vi.hoisted(() => ({
|
|||||||
hasValidAuthSession: vi.fn(),
|
hasValidAuthSession: vi.fn(),
|
||||||
persistAuthSession: vi.fn(),
|
persistAuthSession: vi.fn(),
|
||||||
clearAuthSession: vi.fn(),
|
clearAuthSession: vi.fn(),
|
||||||
|
getSessionToken: vi.fn(),
|
||||||
|
commitSession: vi.fn(),
|
||||||
|
clearSession: vi.fn(),
|
||||||
},
|
},
|
||||||
}));
|
}));
|
||||||
|
|
||||||
@@ -31,6 +34,9 @@ vi.mock('./authSession.js', () => ({
|
|||||||
hasValidAuthSession: authSessionMock.hasValidAuthSession,
|
hasValidAuthSession: authSessionMock.hasValidAuthSession,
|
||||||
persistAuthSession: authSessionMock.persistAuthSession,
|
persistAuthSession: authSessionMock.persistAuthSession,
|
||||||
clearAuthSession: authSessionMock.clearAuthSession,
|
clearAuthSession: authSessionMock.clearAuthSession,
|
||||||
|
getSessionToken: authSessionMock.getSessionToken,
|
||||||
|
commitSession: authSessionMock.commitSession,
|
||||||
|
clearSession: authSessionMock.clearSession,
|
||||||
}));
|
}));
|
||||||
|
|
||||||
vi.mock('./components/SearchModal.js', () => ({
|
vi.mock('./components/SearchModal.js', () => ({
|
||||||
@@ -140,6 +146,9 @@ describe('App mobile sidebar', () => {
|
|||||||
vi.useFakeTimers();
|
vi.useFakeTimers();
|
||||||
apiMock.getEvents.mockResolvedValue([]);
|
apiMock.getEvents.mockResolvedValue([]);
|
||||||
authSessionMock.hasValidAuthSession.mockReturnValue(true);
|
authSessionMock.hasValidAuthSession.mockReturnValue(true);
|
||||||
|
authSessionMock.getSessionToken.mockImplementation((audience: string) =>
|
||||||
|
audience === 'admin' ? 'fake-admin-token' : null,
|
||||||
|
);
|
||||||
});
|
});
|
||||||
|
|
||||||
afterEach(() => {
|
afterEach(() => {
|
||||||
|
|||||||
+33
-366
@@ -1,11 +1,11 @@
|
|||||||
import React, { Suspense, lazy, useState, useEffect, useRef } from 'react';
|
import React, { Suspense, lazy, useState, useEffect, useRef } from 'react';
|
||||||
import { Routes, Route, NavLink, Navigate, useLocation, useNavigate } from 'react-router-dom';
|
import { Routes, Route, NavLink, Navigate, useLocation } from 'react-router-dom';
|
||||||
import { ToastProvider, useToast } from './components/Toast.js';
|
import { ToastProvider, useToast } from './components/Toast.js';
|
||||||
import SearchModal from './components/SearchModal.js';
|
import SearchModal from './components/SearchModal.js';
|
||||||
import NotificationPanel from './components/NotificationPanel.js';
|
import NotificationPanel from './components/NotificationPanel.js';
|
||||||
import TooltipLayer from './components/TooltipLayer.js';
|
import TooltipLayer from './components/TooltipLayer.js';
|
||||||
import { api } from './api.js';
|
import { api } from './api.js';
|
||||||
import { clearAuthSession, clearUserSession, hasValidAuthSession, hasValidUserSession, persistAuthSession, persistUserSession } from './authSession.js';
|
import { clearSession, getSessionToken } from './authSession.js';
|
||||||
import {
|
import {
|
||||||
FIRST_USE_DOC_REMINDER_KEY,
|
FIRST_USE_DOC_REMINDER_KEY,
|
||||||
LEGACY_THEME_STORAGE_KEY,
|
LEGACY_THEME_STORAGE_KEY,
|
||||||
@@ -13,12 +13,13 @@ import {
|
|||||||
USER_PROFILE_STORAGE_KEY,
|
USER_PROFILE_STORAGE_KEY,
|
||||||
} from './appLocalState.js';
|
} from './appLocalState.js';
|
||||||
import { I18nProvider, useI18n } from './i18n.js';
|
import { I18nProvider, useI18n } from './i18n.js';
|
||||||
import { resolveLoginErrorMessage } from './loginError.js';
|
import { SITE_DOCS_URL } from './docsLink.js';
|
||||||
import { SITE_DOCS_URL, SITE_GITHUB_URL } from './docsLink.js';
|
|
||||||
import { useAnimatedVisibility } from './components/useAnimatedVisibility.js';
|
import { useAnimatedVisibility } from './components/useAnimatedVisibility.js';
|
||||||
import { useIsMobile } from './components/useIsMobile.js';
|
import { useIsMobile } from './components/useIsMobile.js';
|
||||||
import { MobileDrawer } from './components/MobileDrawer.js';
|
import { MobileDrawer } from './components/MobileDrawer.js';
|
||||||
import CenteredModal from './components/CenteredModal.js';
|
import CenteredModal from './components/CenteredModal.js';
|
||||||
|
import AdminLoginPage from './pages/auth/AdminLoginPage.js';
|
||||||
|
import UserLoginPage from './pages/auth/UserLoginPage.js';
|
||||||
const Dashboard = lazy(() => import('./pages/Dashboard.js'));
|
const Dashboard = lazy(() => import('./pages/Dashboard.js'));
|
||||||
const Sites = lazy(() => import('./pages/Sites.js'));
|
const Sites = lazy(() => import('./pages/Sites.js'));
|
||||||
const Accounts = lazy(() => import('./pages/Accounts.js'));
|
const Accounts = lazy(() => import('./pages/Accounts.js'));
|
||||||
@@ -126,341 +127,6 @@ function resolveStoredProfile(): UserProfile {
|
|||||||
return { name: '管理员', avatarSeed, avatarStyle: pickDicebearStyle(avatarSeed) };
|
return { name: '管理员', avatarSeed, avatarStyle: pickDicebearStyle(avatarSeed) };
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
type LoginAudience = 'user' | 'admin';
|
|
||||||
type LoginMode = 'token' | 'account';
|
|
||||||
|
|
||||||
export function Login({
|
|
||||||
onLogin,
|
|
||||||
onUserLogin,
|
|
||||||
t,
|
|
||||||
initialAudience = 'user',
|
|
||||||
}: {
|
|
||||||
onLogin: (token: string) => void;
|
|
||||||
onUserLogin?: () => void;
|
|
||||||
t: (text: string) => string;
|
|
||||||
initialAudience?: LoginAudience;
|
|
||||||
}) {
|
|
||||||
const navigate = useNavigate();
|
|
||||||
const [audience, setAudience] = useState<LoginAudience>(initialAudience);
|
|
||||||
const [mode, setMode] = useState<LoginMode>('token');
|
|
||||||
const [token, setToken] = useState('');
|
|
||||||
const [email, setEmail] = useState('');
|
|
||||||
const [password, setPassword] = useState('');
|
|
||||||
const [loading, setLoading] = useState(false);
|
|
||||||
const [error, setError] = useState('');
|
|
||||||
const capabilityRows = [
|
|
||||||
{
|
|
||||||
title: t('统一代理网关'),
|
|
||||||
description: t('一个 Key、一个入口，兼容 OpenAI / Claude 下游格式'),
|
|
||||||
},
|
|
||||||
{
|
|
||||||
title: t('自动模型发现'),
|
|
||||||
description: t('上游新增模型自动出现在模型列表，零配置路由生成'),
|
|
||||||
},
|
|
||||||
{
|
|
||||||
title: t('智能路由引擎'),
|
|
||||||
description: t('按成本、延迟、成功率自动选择最优通道，故障自动转移'),
|
|
||||||
},
|
|
||||||
];
|
|
||||||
const handleTokenLogin = async () => {
|
|
||||||
if (!token) return;
|
|
||||||
setLoading(true);
|
|
||||||
setError('');
|
|
||||||
try {
|
|
||||||
const res = await fetch('/api/settings/auth/info', {
|
|
||||||
headers: { Authorization: `Bearer ${token}` },
|
|
||||||
});
|
|
||||||
if (res.ok) {
|
|
||||||
onLogin(token);
|
|
||||||
} else {
|
|
||||||
let reason = '';
|
|
||||||
try {
|
|
||||||
const text = await res.text();
|
|
||||||
if (text) {
|
|
||||||
try {
|
|
||||||
const parsed = JSON.parse(text) as { message?: unknown; error?: unknown };
|
|
||||||
if (typeof parsed.message === 'string') reason = parsed.message;
|
|
||||||
else if (typeof parsed.error === 'string') reason = parsed.error;
|
|
||||||
else reason = text;
|
|
||||||
} catch {
|
|
||||||
reason = text;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
} catch { }
|
|
||||||
setError(t(resolveLoginErrorMessage(res.status, reason)));
|
|
||||||
setLoading(false);
|
|
||||||
}
|
|
||||||
} catch {
|
|
||||||
setError(t('无法连接到服务器'));
|
|
||||||
setLoading(false);
|
|
||||||
}
|
|
||||||
};
|
|
||||||
const handleAccountLogin = async () => {
|
|
||||||
if (!email.trim() || !password) return;
|
|
||||||
setLoading(true);
|
|
||||||
setError('');
|
|
||||||
try {
|
|
||||||
const res = await fetch('/api/users/login', {
|
|
||||||
method: 'POST',
|
|
||||||
headers: { 'Content-Type': 'application/json' },
|
|
||||||
body: JSON.stringify({ email: email.trim(), password }),
|
|
||||||
});
|
|
||||||
const data = await res.json().catch(() => null);
|
|
||||||
if (!res.ok) {
|
|
||||||
const rawErr = data?.error;
|
|
||||||
const message = data?.message
|
|
||||||
|| (typeof rawErr === 'string' ? rawErr : rawErr?.message)
|
|
||||||
|| t('登录失败，请检查邮箱和密码');
|
|
||||||
setError(t(String(message)));
|
|
||||||
setLoading(false);
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
|
|
||||||
const jwt = String(data?.token || data?.data?.token || '').trim();
|
|
||||||
const role = String(data?.user?.role || data?.data?.user?.role || '').trim();
|
|
||||||
if (!jwt) {
|
|
||||||
setError(t('登录失败，请检查邮箱和密码'));
|
|
||||||
setLoading(false);
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
if (audience === 'admin' && role !== 'admin') {
|
|
||||||
setError(t('该账号不是管理员，无法进入控制台。'));
|
|
||||||
setLoading(false);
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
if (audience === 'admin') {
|
|
||||||
onLogin(jwt);
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
persistUserSession(localStorage, jwt);
|
|
||||||
onUserLogin?.();
|
|
||||||
navigate('/user/dashboard');
|
|
||||||
} catch {
|
|
||||||
setError(t('无法连接到服务器'));
|
|
||||||
setLoading(false);
|
|
||||||
}
|
|
||||||
};
|
|
||||||
const handleLogin = () => {
|
|
||||||
if (audience === 'admin' && mode === 'token') {
|
|
||||||
void handleTokenLogin();
|
|
||||||
} else {
|
|
||||||
void handleAccountLogin();
|
|
||||||
}
|
|
||||||
};
|
|
||||||
const canSubmit = audience === 'admin' && mode === 'token'
|
|
||||||
? !!token.trim()
|
|
||||||
: !!email.trim() && !!password;
|
|
||||||
const activeMode = audience === 'user' ? 'account' : mode;
|
|
||||||
return (
|
|
||||||
<div className="login-shell">
|
|
||||||
<div className="login-surface animate-scale-in">
|
|
||||||
<section className="login-brand-panel login-brand-panel-light">
|
|
||||||
<div className="login-brand-header">
|
|
||||||
<div className="brand-mark-frame brand-mark-frame-hero">
|
|
||||||
<div className="brand-mark-canvas">
|
|
||||||
<img src="/logo.png" alt="BoosAPI" className="login-brand-logo" />
|
|
||||||
</div>
|
|
||||||
</div>
|
|
||||||
<div className="login-brand-summary">
|
|
||||||
<div className="login-brand-name">BoosAPI</div>
|
|
||||||
<div className="login-brand-kicker">{t('中转站的中转站')}</div>
|
|
||||||
</div>
|
|
||||||
</div>
|
|
||||||
<div className="login-brand-copy-block">
|
|
||||||
<p className="login-brand-copy">
|
|
||||||
{t('把分散的 New API / One API / OneHub 等站点聚合成统一网关，自动发现模型、智能路由、成本更优。')}
|
|
||||||
</p>
|
|
||||||
</div>
|
|
||||||
<div className="login-compat-line">{t('兼容 New API / One API / OneHub / DoneHub / Veloera / AnyRouter / Sub2API')}</div>
|
|
||||||
<div className="login-capability-list">
|
|
||||||
{capabilityRows.map((feature, index) => (
|
|
||||||
<div key={feature.title} className="login-capability-row">
|
|
||||||
<div className="login-capability-index">{String(index + 1).padStart(2, '0')}</div>
|
|
||||||
<div className="login-capability-content">
|
|
||||||
<div className="login-capability-title">{feature.title}</div>
|
|
||||||
<p className="login-capability-desc">{feature.description}</p>
|
|
||||||
</div>
|
|
||||||
</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>
|
|
||||||
<section className="login-auth-stage">
|
|
||||||
<div className="login-auth-panel">
|
|
||||||
<div className="login-auth-eyebrow">
|
|
||||||
{audience === 'admin' ? t('管理员入口') : t('用户入口')}
|
|
||||||
</div>
|
|
||||||
<h2 className="login-auth-title">{t('登录')}</h2>
|
|
||||||
<p className="login-auth-copy">
|
|
||||||
{audience === 'user'
|
|
||||||
? t('普通用户使用邮箱和密码进入个人面板。')
|
|
||||||
: activeMode === 'token'
|
|
||||||
? t('请输入管理员令牌后继续。')
|
|
||||||
: t('使用管理员账号和密码进入控制台。')}
|
|
||||||
</p>
|
|
||||||
<div className="login-mode-tabs" role="tablist" aria-label={t('登录入口')}>
|
|
||||||
<button
|
|
||||||
type="button"
|
|
||||||
role="tab"
|
|
||||||
aria-selected={audience === 'user'}
|
|
||||||
className={`login-mode-tab ${audience === 'user' ? 'active' : ''}`}
|
|
||||||
onClick={() => {
|
|
||||||
setAudience('user');
|
|
||||||
setError('');
|
|
||||||
}}
|
|
||||||
>
|
|
||||||
{t('普通用户')}
|
|
||||||
</button>
|
|
||||||
<button
|
|
||||||
type="button"
|
|
||||||
role="tab"
|
|
||||||
aria-selected={audience === 'admin'}
|
|
||||||
className={`login-mode-tab ${audience === 'admin' ? 'active' : ''}`}
|
|
||||||
onClick={() => {
|
|
||||||
setAudience('admin');
|
|
||||||
setError('');
|
|
||||||
}}
|
|
||||||
>
|
|
||||||
{t('管理员')}
|
|
||||||
</button>
|
|
||||||
</div>
|
|
||||||
<div className="login-mode-tabs" role="tablist" aria-label={t('登录方式')}>
|
|
||||||
<button
|
|
||||||
type="button"
|
|
||||||
role="tab"
|
|
||||||
aria-selected={activeMode === 'token'}
|
|
||||||
className={`login-mode-tab ${activeMode === 'token' ? 'active' : ''}`}
|
|
||||||
onClick={() => {
|
|
||||||
setAudience('admin');
|
|
||||||
setMode('token');
|
|
||||||
setError('');
|
|
||||||
}}
|
|
||||||
>
|
|
||||||
{t('Auth Token')}
|
|
||||||
</button>
|
|
||||||
<button
|
|
||||||
type="button"
|
|
||||||
role="tab"
|
|
||||||
aria-selected={activeMode === 'account'}
|
|
||||||
className={`login-mode-tab ${activeMode === 'account' ? 'active' : ''}`}
|
|
||||||
onClick={() => {
|
|
||||||
setMode('account');
|
|
||||||
setError('');
|
|
||||||
}}
|
|
||||||
>
|
|
||||||
{t('账号密码')}
|
|
||||||
</button>
|
|
||||||
</div>
|
|
||||||
{audience === 'admin' && mode === 'token' ? (
|
|
||||||
<>
|
|
||||||
<label className="login-auth-label" htmlFor="admin-token-input">{t('管理员令牌')}</label>
|
|
||||||
<input
|
|
||||||
id="admin-token-input"
|
|
||||||
type="password"
|
|
||||||
placeholder={t('管理员令牌')}
|
|
||||||
value={token}
|
|
||||||
onChange={(e) => {
|
|
||||||
setToken(e.target.value);
|
|
||||||
setError('');
|
|
||||||
}}
|
|
||||||
onKeyDown={(e) => e.key === 'Enter' && handleLogin()}
|
|
||||||
className="login-auth-input"
|
|
||||||
autoComplete="current-password"
|
|
||||||
/>
|
|
||||||
</>
|
|
||||||
) : (
|
|
||||||
<>
|
|
||||||
<label className="login-auth-label" htmlFor="login-email-input">{t('邮箱')}</label>
|
|
||||||
<input
|
|
||||||
id="login-email-input"
|
|
||||||
type="email"
|
|
||||||
placeholder={audience === 'admin' ? t('例如: admin@example.com') : 'you@example.com'}
|
|
||||||
value={email}
|
|
||||||
onChange={(e) => {
|
|
||||||
setEmail(e.target.value);
|
|
||||||
setError('');
|
|
||||||
}}
|
|
||||||
onKeyDown={(e) => e.key === 'Enter' && handleLogin()}
|
|
||||||
className="login-auth-input"
|
|
||||||
autoComplete="username"
|
|
||||||
/>
|
|
||||||
<label className="login-auth-label" htmlFor="login-password-input">{t('密码')}</label>
|
|
||||||
<input
|
|
||||||
id="login-password-input"
|
|
||||||
type="password"
|
|
||||||
placeholder={t('密码')}
|
|
||||||
value={password}
|
|
||||||
onChange={(e) => {
|
|
||||||
setPassword(e.target.value);
|
|
||||||
setError('');
|
|
||||||
}}
|
|
||||||
onKeyDown={(e) => e.key === 'Enter' && handleLogin()}
|
|
||||||
className="login-auth-input"
|
|
||||||
autoComplete="current-password"
|
|
||||||
/>
|
|
||||||
</>
|
|
||||||
)}
|
|
||||||
{error && (
|
|
||||||
<div className="alert alert-error animate-shake" style={{ marginBottom: 12 }}>
|
|
||||||
{error}
|
|
||||||
</div>
|
|
||||||
)}
|
|
||||||
<button
|
|
||||||
onClick={handleLogin}
|
|
||||||
disabled={loading || !canSubmit}
|
|
||||||
className="btn btn-primary login-auth-submit"
|
|
||||||
>
|
|
||||||
{loading ? <><span className="spinner spinner-sm" style={{ borderTopColor: 'white', borderColor: 'rgba(255,255,255,0.3)' }} />{t('验证中...')}</> : t('登录')}
|
|
||||||
</button>
|
|
||||||
<div className="login-auth-note">
|
|
||||||
{audience === 'user'
|
|
||||||
? t('普通用户登录会进入个人面板，可管理自己的 API Key 和用量。')
|
|
||||||
: mode === 'token'
|
|
||||||
? t('仅校验本地服务访问权限，不会把令牌发送到第三方。')
|
|
||||||
: t('账号密码登录会签发管理员会话，不会保存明文密码。')}
|
|
||||||
</div>
|
|
||||||
<div className="login-auth-footer">
|
|
||||||
{audience === 'user' ? (
|
|
||||||
<span>
|
|
||||||
{t('还没有账号？')}
|
|
||||||
<NavLink to="/user/register" style={{ color: 'var(--color-primary)', textDecoration: 'none' }}>
|
|
||||||
{t('注册')}
|
|
||||||
</NavLink>
|
|
||||||
</span>
|
|
||||||
) : (
|
|
||||||
<span>{t('管理员登录后继续。')}</span>
|
|
||||||
)}
|
|
||||||
</div>
|
|
||||||
</div>
|
|
||||||
</section>
|
|
||||||
</div>
|
|
||||||
</div>
|
|
||||||
);
|
|
||||||
}
|
|
||||||
function UserProfileModal({
|
function UserProfileModal({
|
||||||
open,
|
open,
|
||||||
profile,
|
profile,
|
||||||
@@ -638,8 +304,9 @@ function RouteLoadingFallback() {
|
|||||||
}
|
}
|
||||||
function AppShell() {
|
function AppShell() {
|
||||||
const { language, toggleLanguage, t } = useI18n();
|
const { language, toggleLanguage, t } = useI18n();
|
||||||
const [authed, setAuthed] = useState(() => hasValidAuthSession(localStorage));
|
const location = useLocation();
|
||||||
const [userAuthed, setUserAuthed] = useState(() => hasValidUserSession(localStorage));
|
const [authed, setAuthed] = useState(() => getSessionToken('admin') !== null);
|
||||||
|
const [userAuthed, setUserAuthed] = useState(() => getSessionToken('user') !== null);
|
||||||
const [sidebarCollapsed, setSidebarCollapsed] = useState(false);
|
const [sidebarCollapsed, setSidebarCollapsed] = useState(false);
|
||||||
const [showUserMenu, setShowUserMenu] = useState(false);
|
const [showUserMenu, setShowUserMenu] = useState(false);
|
||||||
const userMenuRef = useRef<HTMLDivElement>(null);
|
const userMenuRef = useRef<HTMLDivElement>(null);
|
||||||
@@ -754,7 +421,7 @@ function AppShell() {
|
|||||||
useEffect(() => {
|
useEffect(() => {
|
||||||
if (!authed) return;
|
if (!authed) return;
|
||||||
const check = () => {
|
const check = () => {
|
||||||
if (hasValidAuthSession(localStorage)) return;
|
if (getSessionToken('admin') !== null) return;
|
||||||
setAuthed(false);
|
setAuthed(false);
|
||||||
toast.info(t('会话已过期，请重新登录'));
|
toast.info(t('会话已过期，请重新登录'));
|
||||||
};
|
};
|
||||||
@@ -799,28 +466,38 @@ function AppShell() {
|
|||||||
toast.success(t('个人信息已保存'));
|
toast.success(t('个人信息已保存'));
|
||||||
};
|
};
|
||||||
const isUserPage = location.pathname.startsWith('/user/');
|
const isUserPage = location.pathname.startsWith('/user/');
|
||||||
|
const isAdminLoginPage = location.pathname === '/boos';
|
||||||
const isPublicUserPage = ['/user/login', '/user/register'].includes(location.pathname);
|
const isPublicUserPage = ['/user/login', '/user/register'].includes(location.pathname);
|
||||||
|
// The default `/` route now serves the user-facing login; admins land on `/boos`.
|
||||||
|
const isDefaultLoginRoot = !authed && !userAuthed && location.pathname === '/';
|
||||||
|
|
||||||
if (isUserPage && isPublicUserPage) {
|
if (isUserPage && isPublicUserPage) {
|
||||||
if (userAuthed) return <Navigate to="/user/dashboard" />;
|
if (userAuthed) return <Navigate to="/user/dashboard" />;
|
||||||
if (authed) return <Navigate to="/" />;
|
if (authed) return <Navigate to="/" />;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
if (isAdminLoginPage) {
|
||||||
|
if (authed) return <Navigate to="/" />;
|
||||||
|
return (
|
||||||
|
<AdminLoginPage
|
||||||
|
t={t}
|
||||||
|
onLogin={() => setAuthed(true)}
|
||||||
|
/>
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
if (isDefaultLoginRoot) {
|
||||||
|
return <Navigate to="/user/login" />;
|
||||||
|
}
|
||||||
|
|
||||||
if (!authed && !isUserPage) {
|
if (!authed && !isUserPage) {
|
||||||
if (userAuthed) {
|
if (userAuthed) {
|
||||||
return <Navigate to="/user/dashboard" />;
|
return <Navigate to="/user/dashboard" />;
|
||||||
}
|
}
|
||||||
return (
|
// Anyone hitting an admin-only path without admin auth gets bounced to the
|
||||||
<Login
|
// user login (or dashboard if they are a logged-in user). Admins know to
|
||||||
t={t}
|
// visit `/boos` directly.
|
||||||
initialAudience="admin"
|
return <Navigate to="/user/login" />;
|
||||||
onLogin={(token) => {
|
|
||||||
persistAuthSession(localStorage, token);
|
|
||||||
setAuthed(true);
|
|
||||||
}}
|
|
||||||
onUserLogin={() => setUserAuthed(true)}
|
|
||||||
/>
|
|
||||||
);
|
|
||||||
}
|
}
|
||||||
|
|
||||||
if (isUserPage && !userAuthed && !isPublicUserPage) {
|
if (isUserPage && !userAuthed && !isPublicUserPage) {
|
||||||
@@ -833,17 +510,7 @@ function AppShell() {
|
|||||||
<Routes>
|
<Routes>
|
||||||
<Route
|
<Route
|
||||||
path="/user/login"
|
path="/user/login"
|
||||||
element={(
|
element={<UserLoginPage t={t} onLogin={() => setUserAuthed(true)} />}
|
||||||
<Login
|
|
||||||
t={t}
|
|
||||||
initialAudience="user"
|
|
||||||
onLogin={(token) => {
|
|
||||||
persistAuthSession(localStorage, token);
|
|
||||||
setAuthed(true);
|
|
||||||
}}
|
|
||||||
onUserLogin={() => setUserAuthed(true)}
|
|
||||||
/>
|
|
||||||
)}
|
|
||||||
/>
|
/>
|
||||||
<Route path="/user/register" element={<UserRegister onLogin={() => setUserAuthed(true)} />} />
|
<Route path="/user/register" element={<UserRegister onLogin={() => setUserAuthed(true)} />} />
|
||||||
</Routes>
|
</Routes>
|
||||||
@@ -984,8 +651,8 @@ function AppShell() {
|
|||||||
{t('个人信息')}
|
{t('个人信息')}
|
||||||
</button>
|
</button>
|
||||||
<button onClick={() => {
|
<button onClick={() => {
|
||||||
clearAuthSession(localStorage);
|
clearSession('admin');
|
||||||
clearUserSession(localStorage);
|
clearSession('user');
|
||||||
setAuthed(false);
|
setAuthed(false);
|
||||||
setUserAuthed(false);
|
setUserAuthed(false);
|
||||||
}} className="user-dropdown-item danger">
|
}} className="user-dropdown-item danger">
|
||||||
|
|||||||
+23
-11
@@ -1,4 +1,4 @@
|
|||||||
import { clearAuthSession, getAuthToken, getUserToken, clearUserSession } from "./authSession.js";
|
import { clearSession, getSessionToken } from "./authSession.js";
|
||||||
|
|
||||||
type BufferLike = {
|
type BufferLike = {
|
||||||
from(data: ArrayBuffer): { toString(encoding: "base64"): string };
|
from(data: ArrayBuffer): { toString(encoding: "base64"): string };
|
||||||
@@ -12,10 +12,10 @@ type RequestOptions = RequestInit & {
|
|||||||
};
|
};
|
||||||
|
|
||||||
function requireAuthToken(): string {
|
function requireAuthToken(): string {
|
||||||
const token = getAuthToken(localStorage);
|
const token = getSessionToken('admin', localStorage);
|
||||||
if (!token) {
|
if (!token) {
|
||||||
const hadToken = !!localStorage.getItem("auth_token");
|
const hadToken = !!localStorage.getItem("auth_token");
|
||||||
clearAuthSession(localStorage);
|
clearSession('admin', localStorage);
|
||||||
if (
|
if (
|
||||||
hadToken &&
|
hadToken &&
|
||||||
typeof window !== "undefined" &&
|
typeof window !== "undefined" &&
|
||||||
@@ -127,8 +127,8 @@ async function fetchAuthenticatedResponse(
|
|||||||
headers,
|
headers,
|
||||||
});
|
});
|
||||||
if (res.status === 401 || res.status === 403) {
|
if (res.status === 401 || res.status === 403) {
|
||||||
const hadToken = !!getAuthToken(localStorage);
|
const hadToken = !!getSessionToken('admin', localStorage);
|
||||||
clearAuthSession(localStorage);
|
clearSession('admin', localStorage);
|
||||||
if (
|
if (
|
||||||
hadToken &&
|
hadToken &&
|
||||||
typeof window !== "undefined" &&
|
typeof window !== "undefined" &&
|
||||||
@@ -182,9 +182,9 @@ async function fetchUserResponse(
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
const token = getUserToken(localStorage);
|
const token = getSessionToken('user', localStorage);
|
||||||
if (!token) {
|
if (!token) {
|
||||||
clearUserSession(localStorage);
|
clearSession('user', localStorage);
|
||||||
throw new Error("User session expired");
|
throw new Error("User session expired");
|
||||||
}
|
}
|
||||||
const headers = new Headers(fetchOptions.headers ?? {});
|
const headers = new Headers(fetchOptions.headers ?? {});
|
||||||
@@ -200,7 +200,7 @@ async function fetchUserResponse(
|
|||||||
headers,
|
headers,
|
||||||
});
|
});
|
||||||
if (res.status === 401 || res.status === 403) {
|
if (res.status === 401 || res.status === 403) {
|
||||||
clearUserSession(localStorage);
|
clearSession('user', localStorage);
|
||||||
throw new Error("User session expired");
|
throw new Error("User session expired");
|
||||||
}
|
}
|
||||||
return res;
|
return res;
|
||||||
@@ -325,8 +325,20 @@ function normalizeProxyLogRow(row: any): any {
|
|||||||
routeId: row.routeId ?? row.route_id ?? null,
|
routeId: row.routeId ?? row.route_id ?? null,
|
||||||
channelId: row.channelId ?? row.channel_id ?? null,
|
channelId: row.channelId ?? row.channel_id ?? null,
|
||||||
accountId: row.accountId ?? row.account_id ?? null,
|
accountId: row.accountId ?? row.account_id ?? null,
|
||||||
|
siteId: row.siteId ?? row.site_id ?? null,
|
||||||
|
username: row.username ?? null,
|
||||||
|
siteName: row.siteName ?? row.site_name ?? null,
|
||||||
|
siteUrl: row.siteUrl ?? row.site_url ?? null,
|
||||||
downstreamApiKeyId:
|
downstreamApiKeyId:
|
||||||
row.downstreamApiKeyId ?? row.downstream_api_key_id ?? null,
|
row.downstreamApiKeyId ?? row.downstream_api_key_id ?? null,
|
||||||
|
downstreamKeyId:
|
||||||
|
row.downstreamKeyId ?? row.downstream_key_id ?? row.downstreamApiKeyId ?? row.downstream_api_key_id ?? null,
|
||||||
|
downstreamKeyName:
|
||||||
|
row.downstreamKeyName ?? row.downstream_key_name ?? null,
|
||||||
|
downstreamKeyGroupName:
|
||||||
|
row.downstreamKeyGroupName ?? row.downstream_key_group_name ?? null,
|
||||||
|
downstreamKeyTags:
|
||||||
|
row.downstreamKeyTags ?? parseJsonArrayField(row.downstream_key_tags),
|
||||||
modelRequested: row.modelRequested ?? row.model_requested ?? "",
|
modelRequested: row.modelRequested ?? row.model_requested ?? "",
|
||||||
modelActual: row.modelActual ?? row.model_actual ?? null,
|
modelActual: row.modelActual ?? row.model_actual ?? null,
|
||||||
httpStatus: row.httpStatus ?? row.http_status ?? null,
|
httpStatus: row.httpStatus ?? row.http_status ?? null,
|
||||||
@@ -1381,7 +1393,7 @@ export const api = {
|
|||||||
normalizeProxyDebugTraceDetail,
|
normalizeProxyDebugTraceDetail,
|
||||||
) as Promise<ProxyDebugTraceDetail>,
|
) as Promise<ProxyDebugTraceDetail>,
|
||||||
checkModels: (accountId: number) =>
|
checkModels: (accountId: number) =>
|
||||||
request(`/api/models/check/${accountId}`, { method: "POST" }),
|
request(`/api/accounts/${accountId}/models/refresh`, { method: "POST" }),
|
||||||
getSiteDistribution: () => request("/api/stats/site-distribution"),
|
getSiteDistribution: () => request("/api/stats/site-distribution"),
|
||||||
getSiteTrend: (days = 7) => request(`/api/stats/site-trend?days=${days}`),
|
getSiteTrend: (days = 7) => request(`/api/stats/site-trend?days=${days}`),
|
||||||
getSiteSnapshot: async (days = 7, options?: { refresh?: boolean }) => {
|
getSiteSnapshot: async (days = 7, options?: { refresh?: boolean }) => {
|
||||||
@@ -1795,9 +1807,9 @@ export const api = {
|
|||||||
data: TestChatRequestPayload,
|
data: TestChatRequestPayload,
|
||||||
signal?: AbortSignal,
|
signal?: AbortSignal,
|
||||||
) => {
|
) => {
|
||||||
const token = getAuthToken(localStorage);
|
const token = getSessionToken('admin', localStorage);
|
||||||
if (!token) {
|
if (!token) {
|
||||||
clearAuthSession(localStorage);
|
clearSession('admin', localStorage);
|
||||||
throw new Error("Session expired");
|
throw new Error("Session expired");
|
||||||
}
|
}
|
||||||
return fetch("/api/test/chat/stream", {
|
return fetch("/api/test/chat/stream", {
|
||||||
|
|||||||
@@ -1,4 +1,4 @@
|
|||||||
import { clearAuthSession } from './authSession.js';
|
import { clearSession } from './authSession.js';
|
||||||
|
|
||||||
export const THEME_MODE_STORAGE_KEY = 'theme_mode';
|
export const THEME_MODE_STORAGE_KEY = 'theme_mode';
|
||||||
export const LEGACY_THEME_STORAGE_KEY = 'theme';
|
export const LEGACY_THEME_STORAGE_KEY = 'theme';
|
||||||
@@ -20,7 +20,7 @@ function resolveStorage(storage?: StorageLike | null): StorageLike | null {
|
|||||||
export function clearAppInstallationState(storage?: StorageLike | null): void {
|
export function clearAppInstallationState(storage?: StorageLike | null): void {
|
||||||
const target = resolveStorage(storage);
|
const target = resolveStorage(storage);
|
||||||
if (!target) return;
|
if (!target) return;
|
||||||
clearAuthSession(target as never);
|
clearSession('admin', target as never);
|
||||||
target.removeItem(THEME_MODE_STORAGE_KEY);
|
target.removeItem(THEME_MODE_STORAGE_KEY);
|
||||||
target.removeItem(LEGACY_THEME_STORAGE_KEY);
|
target.removeItem(LEGACY_THEME_STORAGE_KEY);
|
||||||
target.removeItem(USER_PROFILE_STORAGE_KEY);
|
target.removeItem(USER_PROFILE_STORAGE_KEY);
|
||||||
|
|||||||
@@ -2,14 +2,21 @@ import { describe, expect, it } from 'vitest';
|
|||||||
import {
|
import {
|
||||||
AUTH_SESSION_DURATION_MS,
|
AUTH_SESSION_DURATION_MS,
|
||||||
clearAuthSession,
|
clearAuthSession,
|
||||||
|
clearSession,
|
||||||
|
commitSession,
|
||||||
getAuthToken,
|
getAuthToken,
|
||||||
|
getSessionToken,
|
||||||
|
getUserToken,
|
||||||
hasValidAuthSession,
|
hasValidAuthSession,
|
||||||
|
hasValidUserSession,
|
||||||
persistAuthSession,
|
persistAuthSession,
|
||||||
|
persistUserSession,
|
||||||
} from './authSession.js';
|
} from './authSession.js';
|
||||||
|
|
||||||
function createMemoryStorage() {
|
function createMemoryStorage() {
|
||||||
const store = new Map<string, string>();
|
const store = new Map<string, string>();
|
||||||
return {
|
return {
|
||||||
|
raw: store,
|
||||||
getItem(key: string) {
|
getItem(key: string) {
|
||||||
return store.has(key) ? store.get(key)! : null;
|
return store.has(key) ? store.get(key)! : null;
|
||||||
},
|
},
|
||||||
@@ -22,7 +29,7 @@ function createMemoryStorage() {
|
|||||||
};
|
};
|
||||||
}
|
}
|
||||||
|
|
||||||
describe('authSession', () => {
|
describe('authSession (legacy admin API)', () => {
|
||||||
it('persists token with expiration and reads it before expiry', () => {
|
it('persists token with expiration and reads it before expiry', () => {
|
||||||
const storage = createMemoryStorage();
|
const storage = createMemoryStorage();
|
||||||
persistAuthSession(storage, 'token-1', 60_000, 1_000);
|
persistAuthSession(storage, 'token-1', 60_000, 1_000);
|
||||||
@@ -48,3 +55,77 @@ describe('authSession', () => {
|
|||||||
expect(getAuthToken(storage, 2_000)).toBeNull();
|
expect(getAuthToken(storage, 2_000)).toBeNull();
|
||||||
});
|
});
|
||||||
});
|
});
|
||||||
|
|
||||||
|
describe('authSession (legacy user API)', () => {
|
||||||
|
it('persists token with expiration and reads it before expiry', () => {
|
||||||
|
const storage = createMemoryStorage();
|
||||||
|
persistUserSession(storage, 'user-token', 60_000, 1_000);
|
||||||
|
|
||||||
|
expect(getUserToken(storage, 10_000)).toBe('user-token');
|
||||||
|
expect(hasValidUserSession(storage, 10_000)).toBe(true);
|
||||||
|
});
|
||||||
|
|
||||||
|
it('admin and user sessions live in separate slots', () => {
|
||||||
|
const storage = createMemoryStorage();
|
||||||
|
persistAuthSession(storage, 'admin-token', 60_000, 1_000);
|
||||||
|
persistUserSession(storage, 'user-token', 60_000, 1_000);
|
||||||
|
|
||||||
|
expect(getAuthToken(storage, 5_000)).toBe('admin-token');
|
||||||
|
expect(getUserToken(storage, 5_000)).toBe('user-token');
|
||||||
|
|
||||||
|
clearAuthSession(storage);
|
||||||
|
expect(getAuthToken(storage, 5_000)).toBeNull();
|
||||||
|
expect(getUserToken(storage, 5_000)).toBe('user-token');
|
||||||
|
});
|
||||||
|
});
|
||||||
|
|
||||||
|
describe('authSession (commitSession unified API)', () => {
|
||||||
|
it('writes admin tokens to the same slot as persistAuthSession', () => {
|
||||||
|
const storage = createMemoryStorage();
|
||||||
|
commitSession('admin', 'unified-admin', 60_000, storage, 1_000);
|
||||||
|
|
||||||
|
expect(getAuthToken(storage, 5_000)).toBe('unified-admin');
|
||||||
|
expect(getSessionToken('admin', storage, 5_000)).toBe('unified-admin');
|
||||||
|
});
|
||||||
|
|
||||||
|
it('writes user tokens to the same slot as persistUserSession', () => {
|
||||||
|
const storage = createMemoryStorage();
|
||||||
|
commitSession('user', 'unified-user', 60_000, storage, 1_000);
|
||||||
|
|
||||||
|
expect(getUserToken(storage, 5_000)).toBe('unified-user');
|
||||||
|
expect(getSessionToken('user', storage, 5_000)).toBe('unified-user');
|
||||||
|
});
|
||||||
|
|
||||||
|
it('audiences do not bleed across each other', () => {
|
||||||
|
const storage = createMemoryStorage();
|
||||||
|
commitSession('admin', 'a', 60_000, storage, 1_000);
|
||||||
|
commitSession('user', 'u', 60_000, storage, 1_000);
|
||||||
|
|
||||||
|
clearSession('admin', storage);
|
||||||
|
|
||||||
|
expect(getSessionToken('admin', storage, 5_000)).toBeNull();
|
||||||
|
expect(getSessionToken('user', storage, 5_000)).toBe('u');
|
||||||
|
});
|
||||||
|
|
||||||
|
it('empty token clears the slot', () => {
|
||||||
|
const storage = createMemoryStorage();
|
||||||
|
commitSession('admin', 'first', 60_000, storage, 1_000);
|
||||||
|
commitSession('admin', ' ', 60_000, storage, 1_000);
|
||||||
|
|
||||||
|
expect(getSessionToken('admin', storage, 5_000)).toBeNull();
|
||||||
|
});
|
||||||
|
|
||||||
|
it('expired token is purged on read', () => {
|
||||||
|
const storage = createMemoryStorage();
|
||||||
|
commitSession('user', 'expiring', 60_000, storage, 1_000);
|
||||||
|
|
||||||
|
expect(getSessionToken('user', storage, 1_000_000)).toBeNull();
|
||||||
|
});
|
||||||
|
|
||||||
|
it('legacy entries without expiresAt get auto-migrated', () => {
|
||||||
|
const storage = createMemoryStorage();
|
||||||
|
storage.setItem('auth_token', 'legacy-admin');
|
||||||
|
expect(getSessionToken('admin', storage, 1_000)).toBe('legacy-admin');
|
||||||
|
expect(storage.getItem('auth_token_expires_at')).not.toBeNull();
|
||||||
|
});
|
||||||
|
});
|
||||||
|
|||||||
+129
-74
@@ -4,115 +4,170 @@ const USER_JWT_STORAGE_KEY = 'user_jwt_token';
|
|||||||
const USER_JWT_EXPIRES_AT_STORAGE_KEY = 'user_jwt_expires_at';
|
const USER_JWT_EXPIRES_AT_STORAGE_KEY = 'user_jwt_expires_at';
|
||||||
export const AUTH_SESSION_DURATION_MS = 12 * 60 * 60 * 1000;
|
export const AUTH_SESSION_DURATION_MS = 12 * 60 * 60 * 1000;
|
||||||
|
|
||||||
|
export type SessionAudience = 'admin' | 'user';
|
||||||
|
|
||||||
type StorageLike = {
|
type StorageLike = {
|
||||||
getItem: (key: string) => string | null;
|
getItem: (key: string) => string | null;
|
||||||
setItem: (key: string, value: string) => void;
|
setItem: (key: string, value: string) => void;
|
||||||
removeItem: (key: string) => void;
|
removeItem: (key: string) => void;
|
||||||
};
|
};
|
||||||
|
|
||||||
|
type AudienceKeys = {
|
||||||
|
token: string;
|
||||||
|
expiresAt: string;
|
||||||
|
};
|
||||||
|
|
||||||
|
const AUDIENCE_KEYS: Record<SessionAudience, AudienceKeys> = {
|
||||||
|
admin: {
|
||||||
|
token: AUTH_TOKEN_STORAGE_KEY,
|
||||||
|
expiresAt: AUTH_TOKEN_EXPIRES_AT_STORAGE_KEY,
|
||||||
|
},
|
||||||
|
user: {
|
||||||
|
token: USER_JWT_STORAGE_KEY,
|
||||||
|
expiresAt: USER_JWT_EXPIRES_AT_STORAGE_KEY,
|
||||||
|
},
|
||||||
|
};
|
||||||
|
|
||||||
function resolveStorage(storage?: StorageLike | null): StorageLike | null {
|
function resolveStorage(storage?: StorageLike | null): StorageLike | null {
|
||||||
if (storage) return storage;
|
if (storage) return storage;
|
||||||
if (typeof localStorage !== 'undefined') return localStorage;
|
if (typeof localStorage !== 'undefined') return localStorage;
|
||||||
return null;
|
return null;
|
||||||
}
|
}
|
||||||
|
|
||||||
export function clearAuthSession(storage?: StorageLike | null): void {
|
/**
|
||||||
|
* Persist a session token for the given audience. Replaces any existing token.
|
||||||
|
* Empty / whitespace-only tokens clear the session instead.
|
||||||
|
*/
|
||||||
|
export function commitSession(
|
||||||
|
audience: SessionAudience,
|
||||||
|
token: string,
|
||||||
|
ttlMs: number = AUTH_SESSION_DURATION_MS,
|
||||||
|
storage?: StorageLike | null,
|
||||||
|
nowMs: number = Date.now(),
|
||||||
|
): void {
|
||||||
const target = resolveStorage(storage);
|
const target = resolveStorage(storage);
|
||||||
if (!target) return;
|
if (!target) return;
|
||||||
target.removeItem(AUTH_TOKEN_STORAGE_KEY);
|
|
||||||
target.removeItem(AUTH_TOKEN_EXPIRES_AT_STORAGE_KEY);
|
const cleanToken = (token || '').trim();
|
||||||
|
const keys = AUDIENCE_KEYS[audience];
|
||||||
|
if (!cleanToken) {
|
||||||
|
target.removeItem(keys.token);
|
||||||
|
target.removeItem(keys.expiresAt);
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
|
const expiresAt = nowMs + Math.max(1, Math.trunc(ttlMs));
|
||||||
|
target.setItem(keys.token, cleanToken);
|
||||||
|
target.setItem(keys.expiresAt, String(expiresAt));
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/** Drop the session token for the given audience. */
|
||||||
|
export function clearSession(
|
||||||
|
audience: SessionAudience,
|
||||||
|
storage?: StorageLike | null,
|
||||||
|
): void {
|
||||||
|
const target = resolveStorage(storage);
|
||||||
|
if (!target) return;
|
||||||
|
const keys = AUDIENCE_KEYS[audience];
|
||||||
|
target.removeItem(keys.token);
|
||||||
|
target.removeItem(keys.expiresAt);
|
||||||
|
}
|
||||||
|
|
||||||
|
/**
|
||||||
|
* Read the session token for the given audience, returning null if expired or
|
||||||
|
* missing. Auto-migrates legacy entries (token without an expiresAt) to a
|
||||||
|
* fresh TTL.
|
||||||
|
*/
|
||||||
|
export function getSessionToken(
|
||||||
|
audience: SessionAudience,
|
||||||
|
storage?: StorageLike | null,
|
||||||
|
nowMs: number = Date.now(),
|
||||||
|
): string | null {
|
||||||
|
const target = resolveStorage(storage);
|
||||||
|
if (!target) return null;
|
||||||
|
|
||||||
|
const keys = AUDIENCE_KEYS[audience];
|
||||||
|
const token = (target.getItem(keys.token) || '').trim();
|
||||||
|
if (!token) return null;
|
||||||
|
|
||||||
|
const expiresAtRaw = target.getItem(keys.expiresAt);
|
||||||
|
if (!expiresAtRaw) {
|
||||||
|
// Legacy migration: stamp a fresh TTL the first time we see a tokenless entry.
|
||||||
|
commitSession(audience, token, AUTH_SESSION_DURATION_MS, target, nowMs);
|
||||||
|
return token;
|
||||||
|
}
|
||||||
|
|
||||||
|
const expiresAt = Number(expiresAtRaw);
|
||||||
|
if (!Number.isFinite(expiresAt) || expiresAt <= nowMs) {
|
||||||
|
clearSession(audience, target);
|
||||||
|
return null;
|
||||||
|
}
|
||||||
|
|
||||||
|
return token;
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// Legacy thin-shell API. Prefer commitSession / clearSession / getSessionToken.
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
/** @deprecated use {@link clearSession}('admin'). */
|
||||||
|
export function clearAuthSession(storage?: StorageLike | null): void {
|
||||||
|
clearSession('admin', storage);
|
||||||
|
}
|
||||||
|
|
||||||
|
/** @deprecated use {@link commitSession}('admin', ...). */
|
||||||
export function persistAuthSession(
|
export function persistAuthSession(
|
||||||
storage: StorageLike | null | undefined,
|
storage: StorageLike | null | undefined,
|
||||||
token: string,
|
token: string,
|
||||||
ttlMs = AUTH_SESSION_DURATION_MS,
|
ttlMs: number = AUTH_SESSION_DURATION_MS,
|
||||||
nowMs = Date.now(),
|
nowMs: number = Date.now(),
|
||||||
): void {
|
): void {
|
||||||
const target = resolveStorage(storage);
|
commitSession('admin', token, ttlMs, storage, nowMs);
|
||||||
if (!target) return;
|
|
||||||
|
|
||||||
const cleanToken = (token || '').trim();
|
|
||||||
if (!cleanToken) {
|
|
||||||
clearAuthSession(target);
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
|
|
||||||
const expiresAt = nowMs + Math.max(1, Math.trunc(ttlMs));
|
|
||||||
target.setItem(AUTH_TOKEN_STORAGE_KEY, cleanToken);
|
|
||||||
target.setItem(AUTH_TOKEN_EXPIRES_AT_STORAGE_KEY, String(expiresAt));
|
|
||||||
}
|
}
|
||||||
|
|
||||||
export function getAuthToken(storage?: StorageLike | null, nowMs = Date.now()): string | null {
|
/** @deprecated use {@link getSessionToken}('admin'). */
|
||||||
const target = resolveStorage(storage);
|
export function getAuthToken(
|
||||||
if (!target) return null;
|
storage?: StorageLike | null,
|
||||||
|
nowMs: number = Date.now(),
|
||||||
const token = (target.getItem(AUTH_TOKEN_STORAGE_KEY) || '').trim();
|
): string | null {
|
||||||
if (!token) return null;
|
return getSessionToken('admin', storage, nowMs);
|
||||||
|
|
||||||
const expiresAtRaw = target.getItem(AUTH_TOKEN_EXPIRES_AT_STORAGE_KEY);
|
|
||||||
if (!expiresAtRaw) {
|
|
||||||
// Legacy migration: set a default TTL the first time we read an old session.
|
|
||||||
persistAuthSession(target, token, AUTH_SESSION_DURATION_MS, nowMs);
|
|
||||||
return token;
|
|
||||||
}
|
|
||||||
|
|
||||||
const expiresAt = Number(expiresAtRaw);
|
|
||||||
if (!Number.isFinite(expiresAt) || expiresAt <= nowMs) {
|
|
||||||
clearAuthSession(target);
|
|
||||||
return null;
|
|
||||||
}
|
|
||||||
|
|
||||||
return token;
|
|
||||||
}
|
}
|
||||||
|
|
||||||
export function hasValidAuthSession(storage?: StorageLike | null, nowMs = Date.now()): boolean {
|
/** @deprecated derive from {@link getSessionToken}('admin'). */
|
||||||
return !!getAuthToken(storage, nowMs);
|
export function hasValidAuthSession(
|
||||||
|
storage?: StorageLike | null,
|
||||||
|
nowMs: number = Date.now(),
|
||||||
|
): boolean {
|
||||||
|
return !!getSessionToken('admin', storage, nowMs);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/** @deprecated use {@link clearSession}('user'). */
|
||||||
export function clearUserSession(storage?: StorageLike | null): void {
|
export function clearUserSession(storage?: StorageLike | null): void {
|
||||||
const target = resolveStorage(storage);
|
clearSession('user', storage);
|
||||||
if (!target) return;
|
|
||||||
target.removeItem(USER_JWT_STORAGE_KEY);
|
|
||||||
target.removeItem(USER_JWT_EXPIRES_AT_STORAGE_KEY);
|
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/** @deprecated use {@link commitSession}('user', ...). */
|
||||||
export function persistUserSession(
|
export function persistUserSession(
|
||||||
storage: StorageLike | null | undefined,
|
storage: StorageLike | null | undefined,
|
||||||
token: string,
|
token: string,
|
||||||
ttlMs = AUTH_SESSION_DURATION_MS,
|
ttlMs: number = AUTH_SESSION_DURATION_MS,
|
||||||
nowMs = Date.now(),
|
nowMs: number = Date.now(),
|
||||||
): void {
|
): void {
|
||||||
const target = resolveStorage(storage);
|
commitSession('user', token, ttlMs, storage, nowMs);
|
||||||
if (!target) return;
|
|
||||||
const cleanToken = (token || '').trim();
|
|
||||||
if (!cleanToken) {
|
|
||||||
clearUserSession(target);
|
|
||||||
return;
|
|
||||||
}
|
|
||||||
const expiresAt = nowMs + Math.max(1, Math.trunc(ttlMs));
|
|
||||||
target.setItem(USER_JWT_STORAGE_KEY, cleanToken);
|
|
||||||
target.setItem(USER_JWT_EXPIRES_AT_STORAGE_KEY, String(expiresAt));
|
|
||||||
}
|
}
|
||||||
|
|
||||||
export function getUserToken(storage?: StorageLike | null, nowMs = Date.now()): string | null {
|
/** @deprecated use {@link getSessionToken}('user'). */
|
||||||
const target = resolveStorage(storage);
|
export function getUserToken(
|
||||||
if (!target) return null;
|
storage?: StorageLike | null,
|
||||||
const token = (target.getItem(USER_JWT_STORAGE_KEY) || '').trim();
|
nowMs: number = Date.now(),
|
||||||
if (!token) return null;
|
): string | null {
|
||||||
const expiresAtRaw = target.getItem(USER_JWT_EXPIRES_AT_STORAGE_KEY);
|
return getSessionToken('user', storage, nowMs);
|
||||||
if (!expiresAtRaw) {
|
|
||||||
persistUserSession(target, token, AUTH_SESSION_DURATION_MS, nowMs);
|
|
||||||
return token;
|
|
||||||
}
|
|
||||||
const expiresAt = Number(expiresAtRaw);
|
|
||||||
if (!Number.isFinite(expiresAt) || expiresAt <= nowMs) {
|
|
||||||
clearUserSession(target);
|
|
||||||
return null;
|
|
||||||
}
|
|
||||||
return token;
|
|
||||||
}
|
}
|
||||||
|
|
||||||
export function hasValidUserSession(storage?: StorageLike | null, nowMs = Date.now()): boolean {
|
/** @deprecated derive from {@link getSessionToken}('user'). */
|
||||||
return !!getUserToken(storage, nowMs);
|
export function hasValidUserSession(
|
||||||
|
storage?: StorageLike | null,
|
||||||
|
nowMs: number = Date.now(),
|
||||||
|
): boolean {
|
||||||
|
return !!getSessionToken('user', storage, nowMs);
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -2,7 +2,7 @@ import React, { useEffect, useState } from 'react';
|
|||||||
import { createPortal } from 'react-dom';
|
import { createPortal } from 'react-dom';
|
||||||
import { api } from '../api.js';
|
import { api } from '../api.js';
|
||||||
import { useToast } from './Toast.js';
|
import { useToast } from './Toast.js';
|
||||||
import { persistAuthSession } from '../authSession.js';
|
import { commitSession } from '../authSession.js';
|
||||||
import { useAnimatedVisibility } from './useAnimatedVisibility.js';
|
import { useAnimatedVisibility } from './useAnimatedVisibility.js';
|
||||||
|
|
||||||
export default function ChangeKeyModal({ open, onClose }: { open: boolean; onClose: () => void }) {
|
export default function ChangeKeyModal({ open, onClose }: { open: boolean; onClose: () => void }) {
|
||||||
@@ -43,7 +43,7 @@ export default function ChangeKeyModal({ open, onClose }: { open: boolean; onClo
|
|||||||
const res = await api.changeAuthToken(oldToken, newToken);
|
const res = await api.changeAuthToken(oldToken, newToken);
|
||||||
if (res.success) {
|
if (res.success) {
|
||||||
toast.success('Token 已更新，请使用新 Token 重新登录');
|
toast.success('Token 已更新，请使用新 Token 重新登录');
|
||||||
persistAuthSession(localStorage, newToken);
|
commitSession('admin', newToken, undefined, localStorage);
|
||||||
onClose();
|
onClose();
|
||||||
setOldToken('');
|
setOldToken('');
|
||||||
setNewToken('');
|
setNewToken('');
|
||||||
|
|||||||
@@ -4125,6 +4125,70 @@ textarea {
|
|||||||
color: var(--color-text-muted);
|
color: var(--color-text-muted);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
.login-shell[data-variant="user"] .login-surface {
|
||||||
|
width: min(560px, 100%);
|
||||||
|
grid-template-columns: minmax(0, 1fr);
|
||||||
|
gap: 14px;
|
||||||
|
padding: 16px;
|
||||||
|
}
|
||||||
|
|
||||||
|
.login-shell[data-variant="user"] .login-brand-panel {
|
||||||
|
min-height: 0;
|
||||||
|
padding: 28px 24px 18px;
|
||||||
|
border-radius: 26px 26px 18px 18px;
|
||||||
|
}
|
||||||
|
|
||||||
|
.login-shell[data-variant="user"] .login-brand-panel::before {
|
||||||
|
width: 180px;
|
||||||
|
height: 180px;
|
||||||
|
}
|
||||||
|
|
||||||
|
.login-shell[data-variant="user"] .login-brand-panel::after {
|
||||||
|
width: 150px;
|
||||||
|
height: 150px;
|
||||||
|
}
|
||||||
|
|
||||||
|
.login-shell[data-variant="user"] .login-brand-header {
|
||||||
|
justify-content: center;
|
||||||
|
flex-wrap: wrap;
|
||||||
|
}
|
||||||
|
|
||||||
|
.login-shell[data-variant="user"] .login-brand-summary {
|
||||||
|
min-width: 0;
|
||||||
|
}
|
||||||
|
|
||||||
|
.login-shell[data-variant="user"] .brand-mark-frame-hero {
|
||||||
|
width: 76px;
|
||||||
|
height: 76px;
|
||||||
|
border-radius: 24px;
|
||||||
|
}
|
||||||
|
|
||||||
|
.login-shell[data-variant="user"] .brand-mark-canvas {
|
||||||
|
border-radius: 18px;
|
||||||
|
padding: 8px;
|
||||||
|
}
|
||||||
|
|
||||||
|
.login-shell[data-variant="user"] .login-brand-logo {
|
||||||
|
border-radius: 14px;
|
||||||
|
}
|
||||||
|
|
||||||
|
.login-shell[data-variant="user"] .login-brand-name {
|
||||||
|
font-size: 34px;
|
||||||
|
letter-spacing: 0;
|
||||||
|
line-height: 1;
|
||||||
|
}
|
||||||
|
|
||||||
|
.login-shell[data-variant="user"] .login-auth-stage {
|
||||||
|
padding: 0;
|
||||||
|
background: transparent;
|
||||||
|
border-radius: 0;
|
||||||
|
}
|
||||||
|
|
||||||
|
.login-shell[data-variant="user"] .login-auth-panel {
|
||||||
|
width: 100%;
|
||||||
|
max-width: 100%;
|
||||||
|
}
|
||||||
|
|
||||||
[data-theme="dark"] .login-surface,
|
[data-theme="dark"] .login-surface,
|
||||||
[data-theme="dark"] .brand-mark-frame,
|
[data-theme="dark"] .brand-mark-frame,
|
||||||
[data-theme="dark"] .brand-mark-canvas,
|
[data-theme="dark"] .brand-mark-canvas,
|
||||||
|
|||||||
@@ -11,7 +11,35 @@ describe('resolveLoginErrorMessage', () => {
|
|||||||
expect(resolveLoginErrorMessage(401, 'Missing Authorization header')).toBe('登录令牌无效');
|
expect(resolveLoginErrorMessage(401, 'Missing Authorization header')).toBe('登录令牌无效');
|
||||||
});
|
});
|
||||||
|
|
||||||
|
it('returns generic forbidden message for unrelated 403', () => {
|
||||||
|
expect(resolveLoginErrorMessage(403, 'something else')).toBe('没有访问权限');
|
||||||
|
});
|
||||||
|
|
||||||
it('returns server error message for 5xx', () => {
|
it('returns server error message for 5xx', () => {
|
||||||
expect(resolveLoginErrorMessage(500, 'internal error')).toBe('服务端异常，请稍后重试');
|
expect(resolveLoginErrorMessage(500, 'internal error')).toBe('服务端异常，请稍后重试');
|
||||||
|
expect(resolveLoginErrorMessage(503, '')).toBe('服务端异常，请稍后重试');
|
||||||
|
});
|
||||||
|
|
||||||
|
it('routes 400 by reason', () => {
|
||||||
|
expect(resolveLoginErrorMessage(400, 'invalid email format')).toBe('邮箱格式不合法');
|
||||||
|
expect(resolveLoginErrorMessage(400, 'password missing')).toBe('密码不能为空');
|
||||||
|
expect(resolveLoginErrorMessage(400, '')).toBe('请求参数错误');
|
||||||
|
});
|
||||||
|
|
||||||
|
it('routes 409 by reason', () => {
|
||||||
|
expect(resolveLoginErrorMessage(409, 'no_account')).toBe('此邮箱尚未注册');
|
||||||
|
expect(resolveLoginErrorMessage(409, 'NO ACCOUNT')).toBe('此邮箱尚未注册');
|
||||||
|
expect(resolveLoginErrorMessage(409, 'account locked too many tries')).toBe(
|
||||||
|
'账号已被锁定，请稍后重试',
|
||||||
|
);
|
||||||
|
expect(resolveLoginErrorMessage(409, '')).toBe('登录冲突，请稍后再试');
|
||||||
|
});
|
||||||
|
|
||||||
|
it('handles network sentinel', () => {
|
||||||
|
expect(resolveLoginErrorMessage('network', '')).toBe('网络异常，请检查连接');
|
||||||
|
});
|
||||||
|
|
||||||
|
it('falls back to generic message for unmapped statuses', () => {
|
||||||
|
expect(resolveLoginErrorMessage(418, '')).toBe('登录失败，请检查服务状态');
|
||||||
});
|
});
|
||||||
});
|
});
|
||||||
|
|||||||
+37
-2
@@ -1,17 +1,52 @@
|
|||||||
|
export type LoginErrorStatus = number | 'network';
|
||||||
|
|
||||||
function normalizeReason(reason: string): string {
|
function normalizeReason(reason: string): string {
|
||||||
return (reason || '').trim().toLowerCase();
|
return (reason || '').trim().toLowerCase();
|
||||||
}
|
}
|
||||||
|
|
||||||
export function resolveLoginErrorMessage(status: number, reason: string): string {
|
export function resolveLoginErrorMessage(
|
||||||
|
status: LoginErrorStatus,
|
||||||
|
reason: string,
|
||||||
|
): string {
|
||||||
const normalizedReason = normalizeReason(reason);
|
const normalizedReason = normalizeReason(reason);
|
||||||
|
|
||||||
|
if (status === 'network') {
|
||||||
|
return '网络异常，请检查连接';
|
||||||
|
}
|
||||||
|
|
||||||
|
if (status === 400) {
|
||||||
|
if (normalizedReason.includes('invalid email')) {
|
||||||
|
return '邮箱格式不合法';
|
||||||
|
}
|
||||||
|
if (normalizedReason.includes('password')) {
|
||||||
|
return '密码不能为空';
|
||||||
|
}
|
||||||
|
return '请求参数错误';
|
||||||
|
}
|
||||||
|
|
||||||
if (status === 403 && normalizedReason.includes('ip not allowed')) {
|
if (status === 403 && normalizedReason.includes('ip not allowed')) {
|
||||||
return '当前 IP 不在管理白名单中';
|
return '当前 IP 不在管理白名单中';
|
||||||
}
|
}
|
||||||
if (status === 401 || (status === 403 && normalizedReason.includes('invalid token'))) {
|
if (status === 401 || (status === 403 && normalizedReason.includes('invalid token'))) {
|
||||||
return '登录令牌无效';
|
return '登录令牌无效';
|
||||||
}
|
}
|
||||||
if (status >= 500) {
|
if (status === 403) {
|
||||||
|
return '没有访问权限';
|
||||||
|
}
|
||||||
|
|
||||||
|
if (status === 409) {
|
||||||
|
if (normalizedReason.includes('no_account') || normalizedReason.includes('no account')) {
|
||||||
|
return '此邮箱尚未注册';
|
||||||
|
}
|
||||||
|
if (normalizedReason.includes('locked')) {
|
||||||
|
return '账号已被锁定，请稍后重试';
|
||||||
|
}
|
||||||
|
return '登录冲突，请稍后再试';
|
||||||
|
}
|
||||||
|
|
||||||
|
if (typeof status === 'number' && status >= 500) {
|
||||||
return '服务端异常，请稍后重试';
|
return '服务端异常，请稍后重试';
|
||||||
}
|
}
|
||||||
|
|
||||||
return '登录失败，请检查服务状态';
|
return '登录失败，请检查服务状态';
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -307,7 +307,19 @@ export default function Accounts() {
|
|||||||
const verifyFailureHint = buildVerifyFailureHint(verifyResult);
|
const verifyFailureHint = buildVerifyFailureHint(verifyResult);
|
||||||
const addAccountPrereqHint = buildAddAccountPrereqHint(verifyResult);
|
const addAccountPrereqHint = buildAddAccountPrereqHint(verifyResult);
|
||||||
|
|
||||||
const setSegment = (nextSegment: ConnectionsSegment) => {
|
const setSegment = (
|
||||||
|
nextSegment: ConnectionsSegment,
|
||||||
|
event?: React.MouseEvent | React.KeyboardEvent,
|
||||||
|
) => {
|
||||||
|
// Guard against synthetic clicks dispatched by third-party libraries
|
||||||
|
// (e.g. @visactor/vrender's event manager has been observed re-firing
|
||||||
|
// the segmented-button onClick after unrelated interactions, which
|
||||||
|
// silently navigates the page back to the default "session" segment
|
||||||
|
// mid-session). Only honour the call when it was driven by a real
|
||||||
|
// user-initiated DOM event, or when it was invoked programmatically
|
||||||
|
// without an event argument (e.g. test fixtures).
|
||||||
|
if (event && !event.isTrusted) return;
|
||||||
|
if (nextSegment === activeSegment) return;
|
||||||
const params = new URLSearchParams(location.search);
|
const params = new URLSearchParams(location.search);
|
||||||
if (nextSegment === "session") params.delete("segment");
|
if (nextSegment === "session") params.delete("segment");
|
||||||
else params.set("segment", nextSegment);
|
else params.set("segment", nextSegment);
|
||||||
@@ -1437,7 +1449,7 @@ export default function Accounts() {
|
|||||||
<button
|
<button
|
||||||
key={segment.value}
|
key={segment.value}
|
||||||
type="button"
|
type="button"
|
||||||
onClick={() => setSegment(segment.value)}
|
onClick={(event) => setSegment(segment.value, event)}
|
||||||
data-tooltip={segment.tooltip}
|
data-tooltip={segment.tooltip}
|
||||||
data-tooltip-side={segment.tooltipSide}
|
data-tooltip-side={segment.tooltipSide}
|
||||||
data-tooltip-align={segment.tooltipAlign}
|
data-tooltip-align={segment.tooltipAlign}
|
||||||
|
|||||||
@@ -1,7 +1,13 @@
|
|||||||
import React from 'react';
|
import React from 'react';
|
||||||
|
import { useLocation } from 'react-router-dom';
|
||||||
|
import { getSessionToken } from '../authSession.js';
|
||||||
|
|
||||||
export default function ComfyUI() {
|
export default function ComfyUI() {
|
||||||
const [isDark, setIsDark] = React.useState(true);
|
const [isDark, setIsDark] = React.useState(true);
|
||||||
|
const location = useLocation();
|
||||||
|
const audience = location.pathname.startsWith('/user/') ? 'user' : 'admin';
|
||||||
|
const token = getSessionToken(audience, localStorage);
|
||||||
|
const iframeSrc = token ? `/comfy/?token=${encodeURIComponent(token)}` : '';
|
||||||
|
|
||||||
React.useEffect(() => {
|
React.useEffect(() => {
|
||||||
const el = document.documentElement;
|
const el = document.documentElement;
|
||||||
@@ -22,16 +28,28 @@ export default function ComfyUI() {
|
|||||||
background: isDark ? '#1a1a2e' : '#f5f5f5',
|
background: isDark ? '#1a1a2e' : '#f5f5f5',
|
||||||
}}
|
}}
|
||||||
>
|
>
|
||||||
<iframe
|
{iframeSrc ? (
|
||||||
src="/comfyui"
|
<iframe
|
||||||
style={{
|
src={iframeSrc}
|
||||||
width: '100%',
|
style={{
|
||||||
|
width: '100%',
|
||||||
|
height: '100%',
|
||||||
|
border: 'none',
|
||||||
|
}}
|
||||||
|
title="ComfyUI"
|
||||||
|
allow="clipboard-read; clipboard-write"
|
||||||
|
/>
|
||||||
|
) : (
|
||||||
|
<div style={{
|
||||||
height: '100%',
|
height: '100%',
|
||||||
border: 'none',
|
display: 'grid',
|
||||||
}}
|
placeItems: 'center',
|
||||||
title="ComfyUI"
|
color: 'var(--color-text-muted)',
|
||||||
allow="clipboard-read; clipboard-write"
|
fontSize: 14,
|
||||||
/>
|
}}>
|
||||||
|
登录状态已过期，请重新登录。
|
||||||
|
</div>
|
||||||
|
)}
|
||||||
</div>
|
</div>
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|||||||
+216
-10
@@ -1,6 +1,8 @@
|
|||||||
import { useEffect, useMemo, useRef, useState, type CSSProperties } from 'react';
|
import { useEffect, useMemo, useRef, useState, type CSSProperties } from 'react';
|
||||||
import { getAuthToken } from '../authSession.js';
|
import { useLocation } from 'react-router-dom';
|
||||||
|
import { getSessionToken } from '../authSession.js';
|
||||||
import { useToast } from '../components/Toast.js';
|
import { useToast } from '../components/Toast.js';
|
||||||
|
import { filterModelTesterModelNames } from './helpers/modelTesterSession.js';
|
||||||
|
|
||||||
type AgentRole = 'user' | 'assistant' | 'system';
|
type AgentRole = 'user' | 'assistant' | 'system';
|
||||||
|
|
||||||
@@ -54,6 +56,11 @@ type StreamAccumulator = {
|
|||||||
workflowJson: unknown;
|
workflowJson: unknown;
|
||||||
};
|
};
|
||||||
|
|
||||||
|
type StoredAgentSession = {
|
||||||
|
messages?: AgentMessage[];
|
||||||
|
selectedModel?: string;
|
||||||
|
};
|
||||||
|
|
||||||
const WELCOME_MESSAGE: AgentMessage = {
|
const WELCOME_MESSAGE: AgentMessage = {
|
||||||
id: 'welcome',
|
id: 'welcome',
|
||||||
role: 'system',
|
role: 'system',
|
||||||
@@ -71,6 +78,53 @@ function createId(prefix: string) {
|
|||||||
return `${prefix}-${Date.now()}-${Math.random().toString(16).slice(2)}`;
|
return `${prefix}-${Date.now()}-${Math.random().toString(16).slice(2)}`;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
function sessionStorageKey(audience: 'admin' | 'user') {
|
||||||
|
return `boosapi.comfyuiAgent.${audience}.session`;
|
||||||
|
}
|
||||||
|
|
||||||
|
function isValidMessage(value: unknown): value is AgentMessage {
|
||||||
|
if (!value || typeof value !== 'object') return false;
|
||||||
|
const item = value as Partial<AgentMessage>;
|
||||||
|
return (
|
||||||
|
typeof item.id === 'string' &&
|
||||||
|
(item.role === 'user' || item.role === 'assistant' || item.role === 'system') &&
|
||||||
|
typeof item.content === 'string'
|
||||||
|
);
|
||||||
|
}
|
||||||
|
|
||||||
|
function readStoredSession(audience: 'admin' | 'user'): StoredAgentSession {
|
||||||
|
try {
|
||||||
|
const raw = localStorage.getItem(sessionStorageKey(audience));
|
||||||
|
if (!raw) return {};
|
||||||
|
const parsed = JSON.parse(raw) as StoredAgentSession;
|
||||||
|
const messages = Array.isArray(parsed.messages)
|
||||||
|
? parsed.messages.filter(isValidMessage)
|
||||||
|
: undefined;
|
||||||
|
return {
|
||||||
|
messages: messages?.length ? messages : undefined,
|
||||||
|
selectedModel:
|
||||||
|
typeof parsed.selectedModel === 'string' ? parsed.selectedModel : undefined,
|
||||||
|
};
|
||||||
|
} catch {
|
||||||
|
return {};
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
function persistStoredSession(
|
||||||
|
audience: 'admin' | 'user',
|
||||||
|
messages: AgentMessage[],
|
||||||
|
selectedModel: string,
|
||||||
|
) {
|
||||||
|
try {
|
||||||
|
localStorage.setItem(
|
||||||
|
sessionStorageKey(audience),
|
||||||
|
JSON.stringify({ messages, selectedModel }),
|
||||||
|
);
|
||||||
|
} catch {
|
||||||
|
// Ignore quota/private-mode failures; chat still works for the current page.
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
function angleLabel(angle?: string) {
|
function angleLabel(angle?: string) {
|
||||||
if (angle === 'front') return '正面';
|
if (angle === 'front') return '正面';
|
||||||
if (angle === 'back') return '背面';
|
if (angle === 'back') return '背面';
|
||||||
@@ -124,11 +178,25 @@ async function readSseStream(
|
|||||||
|
|
||||||
export default function ComfyUIAgent() {
|
export default function ComfyUIAgent() {
|
||||||
const toast = useToast();
|
const toast = useToast();
|
||||||
const [messages, setMessages] = useState<AgentMessage[]>([WELCOME_MESSAGE]);
|
const location = useLocation();
|
||||||
|
const sessionAudience = location.pathname.startsWith('/user/') ? 'user' : 'admin';
|
||||||
|
const initialStoredSession = useMemo(
|
||||||
|
() => readStoredSession(sessionAudience),
|
||||||
|
[sessionAudience],
|
||||||
|
);
|
||||||
|
const [messages, setMessages] = useState<AgentMessage[]>(
|
||||||
|
() => initialStoredSession.messages || [WELCOME_MESSAGE],
|
||||||
|
);
|
||||||
const [input, setInput] = useState('');
|
const [input, setInput] = useState('');
|
||||||
const [loading, setLoading] = useState(false);
|
const [loading, setLoading] = useState(false);
|
||||||
const [streaming, setStreaming] = useState<StreamAccumulator | null>(null);
|
const [streaming, setStreaming] = useState<StreamAccumulator | null>(null);
|
||||||
const [status, setStatus] = useState<AgentStatus | null>(null);
|
const [status, setStatus] = useState<AgentStatus | null>(null);
|
||||||
|
const [modelOptions, setModelOptions] = useState<string[]>([]);
|
||||||
|
const [selectedModel, setSelectedModel] = useState(
|
||||||
|
() => initialStoredSession.selectedModel || '',
|
||||||
|
);
|
||||||
|
const [modelSearch, setModelSearch] = useState('');
|
||||||
|
const [loadingModels, setLoadingModels] = useState(false);
|
||||||
const chatEndRef = useRef<HTMLDivElement>(null);
|
const chatEndRef = useRef<HTMLDivElement>(null);
|
||||||
const abortRef = useRef<AbortController | null>(null);
|
const abortRef = useRef<AbortController | null>(null);
|
||||||
|
|
||||||
@@ -139,13 +207,38 @@ export default function ComfyUIAgent() {
|
|||||||
}
|
}
|
||||||
return streaming?.workflowJson || null;
|
return streaming?.workflowJson || null;
|
||||||
}, [messages, streaming]);
|
}, [messages, streaming]);
|
||||||
|
const filteredModelOptions = useMemo(
|
||||||
|
() => {
|
||||||
|
const filtered = filterModelTesterModelNames(modelOptions, modelSearch).slice(0, 80);
|
||||||
|
if (selectedModel && !filtered.includes(selectedModel)) {
|
||||||
|
return [selectedModel, ...filtered];
|
||||||
|
}
|
||||||
|
return filtered;
|
||||||
|
},
|
||||||
|
[modelOptions, modelSearch, selectedModel],
|
||||||
|
);
|
||||||
|
|
||||||
useEffect(() => {
|
useEffect(() => {
|
||||||
chatEndRef.current?.scrollIntoView({ behavior: 'smooth', block: 'end' });
|
chatEndRef.current?.scrollIntoView({ behavior: 'smooth', block: 'end' });
|
||||||
}, [messages, streaming?.content]);
|
}, [messages, streaming?.content]);
|
||||||
|
|
||||||
useEffect(() => {
|
useEffect(() => {
|
||||||
const token = getAuthToken(localStorage);
|
const stored = readStoredSession(sessionAudience);
|
||||||
|
setMessages(stored.messages || [WELCOME_MESSAGE]);
|
||||||
|
setSelectedModel(stored.selectedModel || '');
|
||||||
|
setModelSearch('');
|
||||||
|
setStreaming(null);
|
||||||
|
setLoading(false);
|
||||||
|
abortRef.current?.abort();
|
||||||
|
abortRef.current = null;
|
||||||
|
}, [sessionAudience]);
|
||||||
|
|
||||||
|
useEffect(() => {
|
||||||
|
persistStoredSession(sessionAudience, messages, selectedModel);
|
||||||
|
}, [messages, selectedModel, sessionAudience]);
|
||||||
|
|
||||||
|
useEffect(() => {
|
||||||
|
const token = getSessionToken(sessionAudience, localStorage);
|
||||||
if (!token) return;
|
if (!token) return;
|
||||||
void fetch('/api/comfyui-agent/status', {
|
void fetch('/api/comfyui-agent/status', {
|
||||||
headers: { Authorization: `Bearer ${token}` },
|
headers: { Authorization: `Bearer ${token}` },
|
||||||
@@ -153,7 +246,40 @@ export default function ComfyUIAgent() {
|
|||||||
.then((response) => (response.ok ? response.json() : null))
|
.then((response) => (response.ok ? response.json() : null))
|
||||||
.then((payload) => setStatus(payload || null))
|
.then((payload) => setStatus(payload || null))
|
||||||
.catch(() => setStatus(null));
|
.catch(() => setStatus(null));
|
||||||
}, []);
|
}, [sessionAudience]);
|
||||||
|
|
||||||
|
useEffect(() => {
|
||||||
|
const loadRouteModels = async () => {
|
||||||
|
const token = getSessionToken(sessionAudience, localStorage);
|
||||||
|
if (!token) return;
|
||||||
|
setLoadingModels(true);
|
||||||
|
try {
|
||||||
|
const response = await fetch('/api/comfyui-agent/models', {
|
||||||
|
headers: { Authorization: `Bearer ${token}` },
|
||||||
|
});
|
||||||
|
if (!response.ok) throw new Error(`HTTP ${response.status}`);
|
||||||
|
const payload = await response.json();
|
||||||
|
const names = Array.isArray(payload?.models)
|
||||||
|
? payload.models
|
||||||
|
.map((model: unknown) => (typeof model === 'string' ? model.trim() : ''))
|
||||||
|
.filter(Boolean)
|
||||||
|
: [];
|
||||||
|
setModelOptions(names);
|
||||||
|
setSelectedModel((current) => {
|
||||||
|
if (current && names.includes(current)) return current;
|
||||||
|
const configuredModel = String(payload?.defaultModel || status?.model || '').trim();
|
||||||
|
if (configuredModel && names.includes(configuredModel)) return configuredModel;
|
||||||
|
return names[0] || configuredModel;
|
||||||
|
});
|
||||||
|
} catch {
|
||||||
|
toast.error('加载模型路由失败');
|
||||||
|
} finally {
|
||||||
|
setLoadingModels(false);
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
void loadRouteModels();
|
||||||
|
}, [sessionAudience, status?.model, toast]);
|
||||||
|
|
||||||
const exportWorkflow = () => {
|
const exportWorkflow = () => {
|
||||||
const workflow = latestWorkflow || {
|
const workflow = latestWorkflow || {
|
||||||
@@ -250,7 +376,7 @@ export default function ComfyUIAgent() {
|
|||||||
const content = input.trim();
|
const content = input.trim();
|
||||||
if (!content || loading) return;
|
if (!content || loading) return;
|
||||||
|
|
||||||
const token = getAuthToken(localStorage);
|
const token = getSessionToken(sessionAudience, localStorage);
|
||||||
if (!token) {
|
if (!token) {
|
||||||
toast.error('登录状态已过期，请重新登录');
|
toast.error('登录状态已过期，请重新登录');
|
||||||
return;
|
return;
|
||||||
@@ -291,7 +417,10 @@ export default function ComfyUIAgent() {
|
|||||||
'Content-Type': 'application/json',
|
'Content-Type': 'application/json',
|
||||||
Authorization: `Bearer ${token}`,
|
Authorization: `Bearer ${token}`,
|
||||||
},
|
},
|
||||||
body: JSON.stringify({ messages: payloadMessages }),
|
body: JSON.stringify({
|
||||||
|
messages: payloadMessages,
|
||||||
|
model: selectedModel || status?.model,
|
||||||
|
}),
|
||||||
signal: controller.signal,
|
signal: controller.signal,
|
||||||
});
|
});
|
||||||
|
|
||||||
@@ -395,12 +524,45 @@ export default function ComfyUIAgent() {
|
|||||||
<div>
|
<div>
|
||||||
<h1 style={titleStyle}>视频助手</h1>
|
<h1 style={titleStyle}>视频助手</h1>
|
||||||
<div style={statusStyle}>
|
<div style={statusStyle}>
|
||||||
模型 {status?.model || '加载中'} · 上游 {status?.upstreamBase || '检测中'} · ComfyUI {status?.comfyui?.status || 'unknown'}
|
模型 {selectedModel || status?.model || '加载中'} · 上游 {status?.upstreamBase || '检测中'} · ComfyUI {status?.comfyui?.status || 'unknown'}
|
||||||
</div>
|
</div>
|
||||||
</div>
|
</div>
|
||||||
<button type="button" onClick={exportWorkflow} style={secondaryButtonStyle}>
|
<div style={headerActionsStyle}>
|
||||||
导出工作流
|
<label style={modelPickerStyle}>
|
||||||
</button>
|
<span style={modelPickerLabelStyle}>对话模型</span>
|
||||||
|
<input
|
||||||
|
value={modelSearch}
|
||||||
|
onChange={(event) => setModelSearch(event.target.value)}
|
||||||
|
placeholder={loadingModels ? '加载模型...' : '搜索模型路由'}
|
||||||
|
disabled={loading || loadingModels || modelOptions.length === 0}
|
||||||
|
style={modelSearchInputStyle}
|
||||||
|
/>
|
||||||
|
<select
|
||||||
|
value={selectedModel}
|
||||||
|
onChange={(event) => {
|
||||||
|
setSelectedModel(event.target.value);
|
||||||
|
setModelSearch('');
|
||||||
|
}}
|
||||||
|
disabled={loading || loadingModels || modelOptions.length === 0}
|
||||||
|
style={modelSelectStyle}
|
||||||
|
>
|
||||||
|
{modelOptions.length === 0 ? (
|
||||||
|
<option value={status?.model || ''}>
|
||||||
|
{loadingModels ? '模型加载中' : (status?.model || '暂无模型路由')}
|
||||||
|
</option>
|
||||||
|
) : filteredModelOptions.length === 0 ? (
|
||||||
|
<option value={selectedModel}>没有匹配模型</option>
|
||||||
|
) : (
|
||||||
|
filteredModelOptions.map((model) => (
|
||||||
|
<option key={model} value={model}>{model}</option>
|
||||||
|
))
|
||||||
|
)}
|
||||||
|
</select>
|
||||||
|
</label>
|
||||||
|
<button type="button" onClick={exportWorkflow} style={secondaryButtonStyle}>
|
||||||
|
导出工作流
|
||||||
|
</button>
|
||||||
|
</div>
|
||||||
</header>
|
</header>
|
||||||
|
|
||||||
<section style={quickPromptStyle}>
|
<section style={quickPromptStyle}>
|
||||||
@@ -506,6 +668,14 @@ const headerStyle: CSSProperties = {
|
|||||||
background: 'var(--color-surface)',
|
background: 'var(--color-surface)',
|
||||||
};
|
};
|
||||||
|
|
||||||
|
const headerActionsStyle: CSSProperties = {
|
||||||
|
display: 'flex',
|
||||||
|
alignItems: 'center',
|
||||||
|
justifyContent: 'flex-end',
|
||||||
|
gap: 10,
|
||||||
|
flexWrap: 'wrap',
|
||||||
|
};
|
||||||
|
|
||||||
const titleStyle: CSSProperties = {
|
const titleStyle: CSSProperties = {
|
||||||
margin: 0,
|
margin: 0,
|
||||||
fontSize: 18,
|
fontSize: 18,
|
||||||
@@ -518,6 +688,42 @@ const statusStyle: CSSProperties = {
|
|||||||
color: 'var(--color-text-muted)',
|
color: 'var(--color-text-muted)',
|
||||||
};
|
};
|
||||||
|
|
||||||
|
const modelPickerStyle: CSSProperties = {
|
||||||
|
display: 'grid',
|
||||||
|
gridTemplateColumns: 'minmax(130px, 180px) minmax(180px, 260px)',
|
||||||
|
gap: 8,
|
||||||
|
alignItems: 'end',
|
||||||
|
};
|
||||||
|
|
||||||
|
const modelPickerLabelStyle: CSSProperties = {
|
||||||
|
gridColumn: '1 / -1',
|
||||||
|
fontSize: 12,
|
||||||
|
fontWeight: 700,
|
||||||
|
color: 'var(--color-text-muted)',
|
||||||
|
};
|
||||||
|
|
||||||
|
const modelSearchInputStyle: CSSProperties = {
|
||||||
|
minHeight: 36,
|
||||||
|
border: '1px solid var(--color-border)',
|
||||||
|
borderRadius: 8,
|
||||||
|
padding: '7px 10px',
|
||||||
|
background: 'var(--color-bg)',
|
||||||
|
color: 'var(--color-text-primary)',
|
||||||
|
fontSize: 12,
|
||||||
|
outline: 'none',
|
||||||
|
};
|
||||||
|
|
||||||
|
const modelSelectStyle: CSSProperties = {
|
||||||
|
minHeight: 36,
|
||||||
|
border: '1px solid var(--color-border)',
|
||||||
|
borderRadius: 8,
|
||||||
|
padding: '7px 10px',
|
||||||
|
background: 'var(--color-bg)',
|
||||||
|
color: 'var(--color-text-primary)',
|
||||||
|
fontSize: 12,
|
||||||
|
outline: 'none',
|
||||||
|
};
|
||||||
|
|
||||||
const quickPromptStyle: CSSProperties = {
|
const quickPromptStyle: CSSProperties = {
|
||||||
display: 'flex',
|
display: 'flex',
|
||||||
gap: 8,
|
gap: 8,
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
import React, { useCallback, useEffect, useMemo, useRef, useState } from 'react';
|
import React, { useCallback, useEffect, useMemo, useRef, useState } from 'react';
|
||||||
import { api } from '../api.js';
|
import { api } from '../api.js';
|
||||||
import { clearAuthSession, getAuthToken } from '../authSession.js';
|
import { clearSession, getSessionToken } from '../authSession.js';
|
||||||
import {
|
import {
|
||||||
DEBUG_TABS,
|
DEBUG_TABS,
|
||||||
DEFAULT_INPUTS,
|
DEFAULT_INPUTS,
|
||||||
@@ -1555,8 +1555,8 @@ export default function ModelTester() {
|
|||||||
try {
|
try {
|
||||||
const response = await api.proxyTestStream(payload, controller.signal);
|
const response = await api.proxyTestStream(payload, controller.signal);
|
||||||
if (response.status === 401 || response.status === 403) {
|
if (response.status === 401 || response.status === 403) {
|
||||||
const hadToken = Boolean(getAuthToken(localStorage));
|
const hadToken = Boolean(getSessionToken('admin', localStorage));
|
||||||
clearAuthSession(localStorage);
|
clearSession('admin', localStorage);
|
||||||
if (hadToken) window.location.reload();
|
if (hadToken) window.location.reload();
|
||||||
throw new Error('会话已过期');
|
throw new Error('会话已过期');
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -266,14 +266,20 @@ async function copyTextToClipboard(text: string) {
|
|||||||
document.body.removeChild(textarea);
|
document.body.removeChild(textarea);
|
||||||
}
|
}
|
||||||
|
|
||||||
function formatLatency(ms: number) {
|
function formatLatency(ms?: number | null) {
|
||||||
|
if (typeof ms !== "number" || !Number.isFinite(ms)) {
|
||||||
|
return "-";
|
||||||
|
}
|
||||||
if (ms >= 1000) {
|
if (ms >= 1000) {
|
||||||
return `${(ms / 1000).toFixed(ms >= 10000 ? 0 : 1)}s`;
|
return `${(ms / 1000).toFixed(ms >= 10000 ? 0 : 1)}s`;
|
||||||
}
|
}
|
||||||
return `${ms}ms`;
|
return `${ms}ms`;
|
||||||
}
|
}
|
||||||
|
|
||||||
function latencyColor(ms: number) {
|
function latencyColor(ms?: number | null) {
|
||||||
|
if (typeof ms !== "number" || !Number.isFinite(ms)) {
|
||||||
|
return "var(--color-text-muted)";
|
||||||
|
}
|
||||||
if (ms >= 3000) return "var(--color-danger)";
|
if (ms >= 3000) return "var(--color-danger)";
|
||||||
if (ms >= 2000)
|
if (ms >= 2000)
|
||||||
return "color-mix(in srgb, var(--color-warning) 30%, var(--color-danger))";
|
return "color-mix(in srgb, var(--color-warning) 30%, var(--color-danger))";
|
||||||
@@ -285,7 +291,10 @@ function latencyColor(ms: number) {
|
|||||||
return "var(--color-success)";
|
return "var(--color-success)";
|
||||||
}
|
}
|
||||||
|
|
||||||
function latencyBgColor(ms: number) {
|
function latencyBgColor(ms?: number | null) {
|
||||||
|
if (typeof ms !== "number" || !Number.isFinite(ms)) {
|
||||||
|
return "var(--color-bg-muted, transparent)";
|
||||||
|
}
|
||||||
if (ms >= 3000)
|
if (ms >= 3000)
|
||||||
return "color-mix(in srgb, var(--color-danger) 12%, transparent)";
|
return "color-mix(in srgb, var(--color-danger) 12%, transparent)";
|
||||||
if (ms >= 1000)
|
if (ms >= 1000)
|
||||||
|
|||||||
@@ -25,7 +25,7 @@ import {
|
|||||||
resolveRoutingProfilePreset,
|
resolveRoutingProfilePreset,
|
||||||
type RoutingWeights,
|
type RoutingWeights,
|
||||||
} from './helpers/routingProfiles.js';
|
} from './helpers/routingProfiles.js';
|
||||||
import { clearAuthSession } from '../authSession.js';
|
import { clearSession } from '../authSession.js';
|
||||||
import { clearAppInstallationState } from '../appLocalState.js';
|
import { clearAppInstallationState } from '../appLocalState.js';
|
||||||
import { tr } from '../i18n.js';
|
import { tr } from '../i18n.js';
|
||||||
import { generateDownstreamSkKey } from './helpers/generateDownstreamSkKey.js';
|
import { generateDownstreamSkKey } from './helpers/generateDownstreamSkKey.js';
|
||||||
@@ -2615,7 +2615,7 @@ export default function Settings() {
|
|||||||
</button>
|
</button>
|
||||||
<button
|
<button
|
||||||
onClick={() => {
|
onClick={() => {
|
||||||
clearAuthSession(localStorage);
|
clearSession('admin', localStorage);
|
||||||
window.location.reload();
|
window.location.reload();
|
||||||
}}
|
}}
|
||||||
className="btn btn-danger"
|
className="btn btn-danger"
|
||||||
|
|||||||
@@ -6,7 +6,7 @@
|
|||||||
import { useEffect, useMemo, useRef, useState } from 'react';
|
import { useEffect, useMemo, useRef, useState } from 'react';
|
||||||
import { useLocation, useNavigate } from 'react-router-dom';
|
import { useLocation, useNavigate } from 'react-router-dom';
|
||||||
import { api } from '../api.js';
|
import { api } from '../api.js';
|
||||||
import { getAuthToken } from '../authSession.js';
|
import { getSessionToken } from '../authSession.js';
|
||||||
import { getBrand } from '../components/BrandIcon.js';
|
import { getBrand } from '../components/BrandIcon.js';
|
||||||
import CenteredModal from '../components/CenteredModal.js';
|
import CenteredModal from '../components/CenteredModal.js';
|
||||||
import ResponsiveFilterPanel from '../components/ResponsiveFilterPanel.js';
|
import ResponsiveFilterPanel from '../components/ResponsiveFilterPanel.js';
|
||||||
@@ -568,9 +568,14 @@ export default function Sites() {
|
|||||||
|
|
||||||
const handleSaveDisabledModels = async () => {
|
const handleSaveDisabledModels = async () => {
|
||||||
if (!editor || editor.mode !== 'edit') return;
|
if (!editor || editor.mode !== 'edit') return;
|
||||||
|
const siteId = editor.editingSiteId;
|
||||||
|
if (typeof siteId !== 'number' || !Number.isFinite(siteId)) {
|
||||||
|
toast.error('未选中有效站点，无法保存禁用模型');
|
||||||
|
return;
|
||||||
|
}
|
||||||
setDisabledModelsSaving(true);
|
setDisabledModelsSaving(true);
|
||||||
try {
|
try {
|
||||||
await api.updateSiteDisabledModels(editor.editingSiteId, disabledModels);
|
await api.updateSiteDisabledModels(siteId, disabledModels);
|
||||||
try {
|
try {
|
||||||
await api.rebuildRoutes(false, false);
|
await api.rebuildRoutes(false, false);
|
||||||
toast.success('禁用模型列表已保存，路由已重建');
|
toast.success('禁用模型列表已保存，路由已重建');
|
||||||
@@ -621,7 +626,7 @@ export default function Sites() {
|
|||||||
setProbeCompleted(false);
|
setProbeCompleted(false);
|
||||||
|
|
||||||
try {
|
try {
|
||||||
const token = getAuthToken(localStorage);
|
const token = getSessionToken('admin', localStorage);
|
||||||
const params = new URLSearchParams({ scope: probeScope });
|
const params = new URLSearchParams({ scope: probeScope });
|
||||||
if (probeScope === 'single' && probeModel.trim()) params.set('modelName', probeModel.trim());
|
if (probeScope === 'single' && probeModel.trim()) params.set('modelName', probeModel.trim());
|
||||||
const threshold = parseInt(probeLatencyThreshold, 10);
|
const threshold = parseInt(probeLatencyThreshold, 10);
|
||||||
|
|||||||
@@ -1,7 +1,7 @@
|
|||||||
import React, { useState, useEffect, useCallback } from 'react';
|
import React, { useState, useEffect, useCallback } from 'react';
|
||||||
import { useNavigate, Link } from 'react-router-dom';
|
import { useNavigate, Link } from 'react-router-dom';
|
||||||
import { userApi } from '../api.js';
|
import { userApi } from '../api.js';
|
||||||
import { clearUserSession, getUserToken } from '../authSession.js';
|
import { clearSession, getSessionToken } from '../authSession.js';
|
||||||
|
|
||||||
interface UserProfile {
|
interface UserProfile {
|
||||||
id: number;
|
id: number;
|
||||||
@@ -45,7 +45,7 @@ export default function UserDashboard({ onLogout }: { onLogout?: () => void }) {
|
|||||||
const [error, setError] = useState('');
|
const [error, setError] = useState('');
|
||||||
|
|
||||||
const checkAuth = useCallback(() => {
|
const checkAuth = useCallback(() => {
|
||||||
if (!getUserToken(localStorage)) {
|
if (!getSessionToken('user', localStorage)) {
|
||||||
navigate('/user/login');
|
navigate('/user/login');
|
||||||
return false;
|
return false;
|
||||||
}
|
}
|
||||||
@@ -75,7 +75,7 @@ export default function UserDashboard({ onLogout }: { onLogout?: () => void }) {
|
|||||||
useEffect(() => { loadData(); }, [loadData]);
|
useEffect(() => { loadData(); }, [loadData]);
|
||||||
|
|
||||||
const handleLogout = () => {
|
const handleLogout = () => {
|
||||||
clearUserSession(localStorage);
|
clearSession('user', localStorage);
|
||||||
onLogout?.();
|
onLogout?.();
|
||||||
navigate('/user/login', { replace: true });
|
navigate('/user/login', { replace: true });
|
||||||
};
|
};
|
||||||
|
|||||||
+4
-106
@@ -1,109 +1,7 @@
|
|||||||
import React, { useState } from 'react';
|
import UserLoginPage from './auth/UserLoginPage.js';
|
||||||
import { useNavigate, Link } from 'react-router-dom';
|
import { useI18n } from '../i18n.js';
|
||||||
import { persistUserSession } from '../authSession.js';
|
|
||||||
|
|
||||||
export default function UserLogin({ onLogin }: { onLogin?: () => void }) {
|
export default function UserLogin({ onLogin }: { onLogin?: () => void }) {
|
||||||
const navigate = useNavigate();
|
const { t } = useI18n();
|
||||||
const [email, setEmail] = useState('');
|
return <UserLoginPage t={t} onLogin={onLogin} />;
|
||||||
const [password, setPassword] = useState('');
|
|
||||||
const [loading, setLoading] = useState(false);
|
|
||||||
const [error, setError] = useState('');
|
|
||||||
|
|
||||||
const handleSubmit = async (e: React.FormEvent) => {
|
|
||||||
e.preventDefault();
|
|
||||||
if (!email.trim() || !password) return;
|
|
||||||
setLoading(true);
|
|
||||||
setError('');
|
|
||||||
try {
|
|
||||||
const res = await fetch('/api/users/login', {
|
|
||||||
method: 'POST',
|
|
||||||
headers: { 'Content-Type': 'application/json' },
|
|
||||||
body: JSON.stringify({ email: email.trim(), password }),
|
|
||||||
});
|
|
||||||
if (!res.ok) {
|
|
||||||
const data = await res.json().catch(() => ({}));
|
|
||||||
const errMsg = typeof data?.error === 'string' ? data.error : (data?.error?.message || data?.message || `登录失败 (${res.status})`);
|
|
||||||
throw new Error(errMsg);
|
|
||||||
}
|
|
||||||
const result = await res.json();
|
|
||||||
persistUserSession(localStorage, result.token);
|
|
||||||
onLogin?.();
|
|
||||||
navigate('/user/dashboard');
|
|
||||||
} catch (err: any) {
|
|
||||||
setError(err?.message || '登录失败，请检查邮箱和密码');
|
|
||||||
} finally {
|
|
||||||
setLoading(false);
|
|
||||||
}
|
|
||||||
};
|
|
||||||
|
|
||||||
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" />
|
|
||||||
</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">使用邮箱和密码登录你的账号。</p>
|
|
||||||
<form onSubmit={handleSubmit}>
|
|
||||||
<label className="login-auth-label" htmlFor="user-email-input">邮箱</label>
|
|
||||||
<input
|
|
||||||
id="user-email-input"
|
|
||||||
type="email"
|
|
||||||
placeholder="your@email.com"
|
|
||||||
value={email}
|
|
||||||
onChange={(e) => { setEmail(e.target.value); setError(''); }}
|
|
||||||
className="login-auth-input"
|
|
||||||
autoComplete="email"
|
|
||||||
required
|
|
||||||
/>
|
|
||||||
<label className="login-auth-label" htmlFor="user-password-input" style={{ marginTop: 12 }}>密码</label>
|
|
||||||
<input
|
|
||||||
id="user-password-input"
|
|
||||||
type="password"
|
|
||||||
placeholder="••••••••"
|
|
||||||
value={password}
|
|
||||||
onChange={(e) => { setPassword(e.target.value); setError(''); }}
|
|
||||||
className="login-auth-input"
|
|
||||||
autoComplete="current-password"
|
|
||||||
required
|
|
||||||
/>
|
|
||||||
{error && (
|
|
||||||
<div className="alert alert-error animate-shake" style={{ marginBottom: 12, marginTop: 8 }}>
|
|
||||||
{error}
|
|
||||||
</div>
|
|
||||||
)}
|
|
||||||
<button
|
|
||||||
type="submit"
|
|
||||||
disabled={loading || !email.trim() || !password}
|
|
||||||
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)' }} /> 登录中...</> : '登录'}
|
|
||||||
</button>
|
|
||||||
</form>
|
|
||||||
<div className="login-auth-footer" style={{ flexDirection: 'column', gap: 8, marginTop: 16 }}>
|
|
||||||
<span>
|
|
||||||
还没有账号？<Link to="/user/register" 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>
|
|
||||||
);
|
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -1,6 +1,6 @@
|
|||||||
import React, { useState } from 'react';
|
import React, { useState } from 'react';
|
||||||
import { useNavigate, Link } from 'react-router-dom';
|
import { useNavigate, Link } from 'react-router-dom';
|
||||||
import { persistUserSession } from '../authSession.js';
|
import { commitSession } from '../authSession.js';
|
||||||
|
|
||||||
export default function UserRegister({ onLogin }: { onLogin?: () => void }) {
|
export default function UserRegister({ onLogin }: { onLogin?: () => void }) {
|
||||||
const navigate = useNavigate();
|
const navigate = useNavigate();
|
||||||
@@ -35,7 +35,7 @@ export default function UserRegister({ onLogin }: { onLogin?: () => void }) {
|
|||||||
throw new Error(data?.error || `注册失败 (${res.status})`);
|
throw new Error(data?.error || `注册失败 (${res.status})`);
|
||||||
}
|
}
|
||||||
const result = await res.json();
|
const result = await res.json();
|
||||||
persistUserSession(localStorage, result.token);
|
commitSession('user', result.token, undefined, localStorage);
|
||||||
onLogin?.();
|
onLogin?.();
|
||||||
navigate('/user/dashboard');
|
navigate('/user/dashboard');
|
||||||
} catch (err: any) {
|
} catch (err: any) {
|
||||||
|
|||||||
@@ -0,0 +1,111 @@
|
|||||||
|
import React, { useState } from 'react';
|
||||||
|
import { commitSession } from '../../authSession.js';
|
||||||
|
import { resolveLoginErrorMessage } from '../../loginError.js';
|
||||||
|
import LoginShell from './LoginShell.js';
|
||||||
|
|
||||||
|
interface AdminLoginPageProps {
|
||||||
|
t: (text: string) => string;
|
||||||
|
/** Called with the admin auth token when login succeeds. */
|
||||||
|
onLogin: (token: string) => void;
|
||||||
|
}
|
||||||
|
|
||||||
|
async function readErrorReason(res: Response): Promise<string> {
|
||||||
|
try {
|
||||||
|
const text = await res.text();
|
||||||
|
if (!text) return '';
|
||||||
|
try {
|
||||||
|
const parsed = JSON.parse(text) as { message?: unknown; error?: unknown };
|
||||||
|
if (typeof parsed.message === 'string') return parsed.message;
|
||||||
|
if (typeof parsed.error === 'string') return parsed.error;
|
||||||
|
const nestedMessage =
|
||||||
|
parsed.error && typeof parsed.error === 'object'
|
||||||
|
? (parsed.error as { message?: unknown }).message
|
||||||
|
: undefined;
|
||||||
|
if (typeof nestedMessage === 'string') return nestedMessage;
|
||||||
|
return text;
|
||||||
|
} catch {
|
||||||
|
return text;
|
||||||
|
}
|
||||||
|
} catch {
|
||||||
|
return '';
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
export default function AdminLoginPage({ t, onLogin }: AdminLoginPageProps) {
|
||||||
|
const [token, setToken] = useState('');
|
||||||
|
const [loading, setLoading] = useState(false);
|
||||||
|
const [error, setError] = useState('');
|
||||||
|
|
||||||
|
const handleTokenLogin = async () => {
|
||||||
|
if (!token.trim()) return;
|
||||||
|
setLoading(true);
|
||||||
|
setError('');
|
||||||
|
try {
|
||||||
|
const res = await fetch('/api/settings/auth/info', {
|
||||||
|
headers: { Authorization: `Bearer ${token}` },
|
||||||
|
});
|
||||||
|
if (res.ok) {
|
||||||
|
commitSession('admin', token);
|
||||||
|
onLogin(token);
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
const reason = await readErrorReason(res);
|
||||||
|
setError(t(resolveLoginErrorMessage(res.status, reason)));
|
||||||
|
} catch {
|
||||||
|
setError(t(resolveLoginErrorMessage('network', '')));
|
||||||
|
} finally {
|
||||||
|
setLoading(false);
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
const canSubmit = !!token.trim();
|
||||||
|
|
||||||
|
return (
|
||||||
|
<LoginShell
|
||||||
|
t={t}
|
||||||
|
variant="admin"
|
||||||
|
eyebrow={t('管理员入口')}
|
||||||
|
title={t('登录')}
|
||||||
|
description={t('请输入管理员令牌后继续。')}
|
||||||
|
error={error}
|
||||||
|
footer={<span>{t('管理员登录后继续。')}</span>}
|
||||||
|
>
|
||||||
|
<label className="login-auth-label" htmlFor="admin-token-input">
|
||||||
|
{t('管理员令牌')}
|
||||||
|
</label>
|
||||||
|
<input
|
||||||
|
id="admin-token-input"
|
||||||
|
type="password"
|
||||||
|
placeholder={t('管理员令牌')}
|
||||||
|
value={token}
|
||||||
|
onChange={(e) => {
|
||||||
|
setToken(e.target.value);
|
||||||
|
setError('');
|
||||||
|
}}
|
||||||
|
onKeyDown={(e) => e.key === 'Enter' && handleTokenLogin()}
|
||||||
|
className="login-auth-input"
|
||||||
|
autoComplete="current-password"
|
||||||
|
/>
|
||||||
|
<button
|
||||||
|
onClick={handleTokenLogin}
|
||||||
|
disabled={loading || !canSubmit}
|
||||||
|
className="btn btn-primary login-auth-submit"
|
||||||
|
>
|
||||||
|
{loading ? (
|
||||||
|
<>
|
||||||
|
<span
|
||||||
|
className="spinner spinner-sm"
|
||||||
|
style={{ borderTopColor: 'white', borderColor: 'rgba(255,255,255,0.3)' }}
|
||||||
|
/>
|
||||||
|
{t('验证中...')}
|
||||||
|
</>
|
||||||
|
) : (
|
||||||
|
t('登录')
|
||||||
|
)}
|
||||||
|
</button>
|
||||||
|
<div className="login-auth-note">
|
||||||
|
{t('仅校验本地服务访问权限，不会把令牌发送到第三方。')}
|
||||||
|
</div>
|
||||||
|
</LoginShell>
|
||||||
|
);
|
||||||
|
}
|
||||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user