boos/BoosAPI
2
0
Fork 2
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor: ComfyUI Agent with structured parsing + image gen + TTS

Browse Source 
- New system prompt with 【CHARACTER】【SCENE】【TTS】【WORKFLOW】 markers
- Backend parses structured sections, auto-calls gpt-image-2 for character/scene images
- Auto-calls TTS API for dialogue lines
- Case-insensitive field parsing, robust whitespace handling
- Static file serving with SPA fallback for frontend
- No Node.js runtime dependency for serving
This commit is contained in:
boos
2026-05-18 04:02:02 +08:00
parent 56391f2193
commit 725f44699a
14 changed files with 1621 additions and 69 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+13 -10
View File
@@ -2,22 +2,23 @@
# BoosAPI
**AI 聚合网关 + 内置 ComfyUI — 将分散的中转站整合为一个统一入口**
**AI 聚合网关 + ComfyUI Agent — 将分散的中转站整合为一个统一入口**
<p>
把你各处注册的 New API / One API / OneHub / DoneHub / Veloera / AnyRouter / Sub2API 等站点，
<br>
汇聚成 <strong>一个 API Key、一个入口</strong>，自动发现模型、智能路由、成本最优。
<br>
<strong>内置 ComfyUI 对话助手</strong>，自然语言生成图像/视频工作流，开箱即用。
<strong>内置 ComfyUI 视频助手</strong> — 自然语言描述脚本，自动生成角色、场景、TTS、ComfyUI 工作流。
<br>
<strong>SaaS 多租户架构</strong>，组织管理、成员邀请、套餐限制，开箱即用。
</p>
<p align="center">
<img alt="Node.js" src="https://img.shields.io/badge/Node.js-22.15%2B-339933?logo=node.js&style=flat">
<img alt="Rust" src="https://img.shields.io/badge/Rust-1.85%2B-000000?logo=rust&style=flat">
<img alt="TypeScript" src="https://img.shields.io/badge/TypeScript-5.x-3178C6?logo=typescript&style=flat">
<img alt="License" src="https://img.shields.io/badge/license-MIT-brightgreen?style=flat">
<img alt="React" src="https://img.shields.io/badge/React-19-61DAFB?logo=react&style=flat">
</p>
<p align="center">
@@ -58,14 +59,16 @@
## ✨ 核心功能
### 🎨 内置 ComfyUI 对话助手
### 🎨 内置 ComfyUI 视频助手（Rust）
- **自然语言生成工作流**：描述你想要的效果，AI 自动构建 ComfyUI 工作流
- **定向视图生成**：前/后/左/右 多角度图像生成
- **场景/视频生成**：文字描述直接生成视频内容
- **SSE 实时流式响应**：生成进度实时反馈
- 侧边栏「对话助手」入口，与聚合网关同一界面
- 自动对接 ComfyUI API（`:8188`），无需额外配置
- **自然语言描述脚本**：说想法 → AI 自动提取角色、场景、对白
- **自动生成四视图**：每个角色正/背/左/右 4 角度，调用 gpt-image-2 生成图片
- **场景生成**：自动为每个场景生成配图
- **TTS 语音合成**：对话台词自动转语音
- **ComfyUI 工作流导出**：完整 workflow JSON，可直接导入 ComfyUI
- **Rust 后端**（boosapi-rs）：高性能 SSE 流式响应，零 Node.js 运行时依赖
- 侧边栏「视频助手」入口，与聚合网关同一界面
- 自动对接上游 API（`api.boos.lat/v1`），无需额外配置
### 🌐 统一代理网关
boosapi-rs/Cargo.toml
+1
View File
@@ -24,6 +24,7 @@ sqlx = { version = "0.8", features = [
# HTTP client
reqwest = { version = "0.12", features = ["stream", "json"] }
lettre = { version = "0.11", default-features = false, features = ["builder", "hostname", "pool", "rustls-tls", "smtp-transport"] }
# Serialization
serde = { version = "1", features = ["derive"] }
boosapi-rs/boosapi.db
BIN
View File
Binary file not shown.
boosapi-rs/comfy/boosapi_comfy_nodes.py
+218
View File
@@ -0,0 +1,218 @@
"""
BoosAPI ComfyUI Custom Nodes
让 ComfyUI workflow 调用 BoosAPI 的 AI 模型（图像生成、LLM 等）。
配置：通过环境变量设置 BoosAPI 地址，或使用默认 http://127.0.0.1:4000/v1
环境变量：
BOOSAPI_BASE — BoosAPI 地址（默认 http://127.0.0.1:4000）
BOOSAPI_KEY — API Key（默认 sk-boos4721）
"""
import json
import os
import requests
# ── Config ───────────────────────────────────────────────────────────────
BOOSAPI_BASE = os.environ.get("BOOSAPI_BASE", "http://127.0.0.1:4000")
BOOSAPI_KEY = os.environ.get("BOOSAPI_KEY", "sk-boos4721")
HEADERS = {
"Authorization": f"Bearer {BOOSAPI_KEY}",
"Content-Type": "application/json",
}
# ── Node Mappings ────────────────────────────────────────────────────────
NODE_CLASS_MAPPINGS = {}
NODE_DISPLAY_NAME_MAPPINGS = {}
# ═════════════════════════════════════════════════════════════════════════
# Node: BoosAPIImageGen — 图像生成 (OpenAI DALL-E / 兼容格式)
# ═════════════════════════════════════════════════════════════════════════
class BoosAPIImageGen:
"""调用 BoosAPI /v1/images/generations 生成图像"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"prompt": ("STRING", {"multiline": True, "default": ""}),
"model": (["dall-e-3", "dall-e-2", "gpt-image-2"], {"default": "gpt-image-2"}),
"size": (["1024x1024", "1792x1024", "1024x1792"], {"default": "1024x1024"}),
"quality": (["standard", "hd"], {"default": "standard"}),
"n": ("INT", {"default": 1, "min": 1, "max": 4}),
},
"optional": {
"style": (["vivid", "natural"], {"default": "vivid"}),
},
}
RETURN_TYPES = ("IMAGE", "STRING")
RETURN_NAMES = ("image", "metadata")
FUNCTION = "generate"
CATEGORY = "BoosAPI"
def generate(self, prompt, model, size, quality, n, style="vivid"):
url = f"{BOOSAPI_BASE}/v1/images/generations"
payload = {
"model": model,
"prompt": prompt,
"n": n,
"size": size,
"quality": quality,
"style": style,
"response_format": "b64_json",
}
try:
resp = requests.post(url, headers=HEADERS, json=payload, timeout=120)
resp.raise_for_status()
data = resp.json()
images = data.get("data", [])
# Return first image as b64
from io import BytesIO
import base64
import torch
from PIL import Image
b64_data = images[0]["b64_json"]
img_bytes = base64.b64decode(b64_data)
pil_img = Image.open(BytesIO(img_bytes)).convert("RGB")
img_tensor = torch.from_numpy(
__import__("numpy").array(pil_img).astype("float32") / 255.0
).unsqueeze(0)
metadata = json.dumps({
"model": model,
"size": size,
"revised_prompt": data.get("data", [{}])[0].get("revised_prompt", ""),
})
return (img_tensor, metadata)
except Exception as e:
print(f"[BoosAPI] ImageGen failed: {e}")
raise
# ═════════════════════════════════════════════════════════════════════════
# Node: BoosAPILLM — LLM Chat (流式文本生成)
# ═════════════════════════════════════════════════════════════════════════
class BoosAPILLM:
"""调用 BoosAPI /v1/chat/completions 生成文本"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"prompt": ("STRING", {"multiline": True, "default": ""}),
"model": ("STRING", {"default": "gpt-5.5"}),
"system_prompt": ("STRING", {"multiline": True, "default": ""}),
"temperature": ("FLOAT", {"default": 0.7, "min": 0.0, "max": 2.0}),
"max_tokens": ("INT", {"default": 2048, "min": 1, "max": 16384}),
},
}
RETURN_TYPES = ("STRING",)
RETURN_NAMES = ("text",)
FUNCTION = "chat"
CATEGORY = "BoosAPI"
def chat(self, prompt, model, system_prompt, temperature, max_tokens):
url = f"{BOOSAPI_BASE}/v1/chat/completions"
messages = []
if system_prompt:
messages.append({"role": "system", "content": system_prompt})
messages.append({"role": "user", "content": prompt})
payload = {
"model": model,
"messages": messages,
"temperature": temperature,
"max_tokens": max_tokens,
"stream": False,
}
try:
resp = requests.post(url, headers=HEADERS, json=payload, timeout=120)
resp.raise_for_status()
data = resp.json()
text = data["choices"][0]["message"]["content"]
return (text,)
except Exception as e:
print(f"[BoosAPI] LLM chat failed: {e}")
raise
# ═════════════════════════════════════════════════════════════════════════
# Node: BoosAPITextToImage — 文生图 (ComfyUI 原生格式)
# ═════════════════════════════════════════════════════════════════════════
class BoosAPITextToImage:
"""输入文本提示词，通过 BoosAPI 生成图像，输出 ComfyUI IMAGE 格式"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"positive": ("STRING", {"multiline": True, "default": ""}),
"model": (["gpt-image-2", "dall-e-3"], {"default": "gpt-image-2"}),
},
"optional": {
"negative": ("STRING", {"multiline": True, "default": ""}),
"size": (["1024x1024", "1792x1024", "1024x1792"], {"default": "1024x1024"}),
},
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "generate"
CATEGORY = "BoosAPI"
def generate(self, positive, model, negative="", size="1024x1024"):
# Build prompt with negative guidance
prompt = positive
if negative:
prompt = f"{positive}\n避免: {negative}"
url = f"{BOOSAPI_BASE}/v1/images/generations"
payload = {
"model": model,
"prompt": prompt,
"n": 1,
"size": size,
"response_format": "b64_json",
}
try:
resp = requests.post(url, headers=HEADERS, json=payload, timeout=120)
resp.raise_for_status()
data = resp.json()
import base64
from io import BytesIO
import torch
from PIL import Image
import numpy as np
b64 = data["data"][0]["b64_json"]
img = Image.open(BytesIO(base64.b64decode(b64))).convert("RGB")
tensor = torch.from_numpy(np.array(img).astype("float32") / 255.0).unsqueeze(0)
return (tensor,)
except Exception as e:
print(f"[BoosAPI] T2I failed: {e}")
raise
# ═════════════════════════════════════════════════════════════════════════
# Register Nodes
# ═════════════════════════════════════════════════════════════════════════
NODE_CLASS_MAPPINGS = {
"BoosAPIImageGen": BoosAPIImageGen,
"BoosAPILLM": BoosAPILLM,
"BoosAPITextToImage": BoosAPITextToImage,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"BoosAPIImageGen": "BoosAPI 图像生成",
"BoosAPILLM": "BoosAPI LLM 对话",
"BoosAPITextToImage": "BoosAPI 文生图",
}
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
boosapi-rs/config.toml
+4 -4
View File
@@ -1,17 +1,17 @@
[server]
host = "0.0.0.0"
port = 3002
port = 4000
log_level = "info"
cors_origins = ["*"]
[db]
url = "sqlite://../data/hub.db?mode=rwc"
url = "sqlite:boosapi.db?mode=rwc"
max_connections = 10
ssl = false
[auth]
jwt_secret = "dev-jwt-secret-change-in-production"
jwt_expiry_hours = 168
jwt_secret = "boosapi-dev-secret-change-in-production"
jwt_expiry_hours = 72
auth_token = "change-me-admin-token"
proxy_token = "sk-boos4721"
boosapi-rs/migrations/20260518000001_usage_logs.sql
+37
View File
@@ -0,0 +1,37 @@
-- Usage tracking tables migrated from the Node.js backend.
CREATE TABLE IF NOT EXISTS usage_logs (
id INTEGER PRIMARY KEY AUTOINCREMENT,
downstream_key_id INTEGER,
model TEXT,
prompt_tokens INTEGER DEFAULT 0,
completion_tokens INTEGER DEFAULT 0,
total_tokens INTEGER DEFAULT 0,
cost REAL DEFAULT 0.0,
created_at TEXT DEFAULT (datetime('now')),
FOREIGN KEY (downstream_key_id) REFERENCES downstream_api_keys(id) ON DELETE SET NULL
);
CREATE INDEX IF NOT EXISTS usage_logs_downstream_key_created_idx
ON usage_logs (downstream_key_id, created_at);
CREATE INDEX IF NOT EXISTS usage_logs_created_idx
ON usage_logs (created_at);
CREATE TABLE IF NOT EXISTS usage_daily_snapshots (
id INTEGER PRIMARY KEY AUTOINCREMENT,
local_day TEXT NOT NULL,
downstream_key_id INTEGER,
request_count INTEGER NOT NULL DEFAULT 0,
prompt_tokens INTEGER NOT NULL DEFAULT 0,
completion_tokens INTEGER NOT NULL DEFAULT 0,
total_tokens INTEGER NOT NULL DEFAULT 0,
total_cost REAL NOT NULL DEFAULT 0.0,
created_at TEXT DEFAULT (datetime('now')),
updated_at TEXT DEFAULT (datetime('now')),
UNIQUE(local_day, downstream_key_id),
FOREIGN KEY (downstream_key_id) REFERENCES downstream_api_keys(id) ON DELETE SET NULL
);
CREATE INDEX IF NOT EXISTS usage_daily_snapshots_day_idx
ON usage_daily_snapshots (local_day);
boosapi-rs/src/routes/comfyui_agent.rs
+465 -49
View File
@@ -7,6 +7,7 @@ use axum::{
Json, Router,
};
use futures::stream::Stream;
use regex::Regex;
use reqwest::Client;
use serde_json::{json, Value};
use std::convert::Infallible;
@@ -14,35 +15,268 @@ use tokio::sync::mpsc::{unbounded_channel, UnboundedSender};
use tokio_stream::{wrappers::UnboundedReceiverStream, StreamExt};
use tracing::info;
const SYSTEM_PROMPT: &str = r#"你是 BoosAPI 的 AI 视频生成助手，帮助用户通过 ComfyUI 制作 AI 视频。
const SYSTEM_PROMPT: &str = r#"你是 BoosAPI 的 AI 视频生成助手，帮用户通过 AI 制作视频。
## 能力
1. **分析脚本** — 提取角色、场景、对话
2. **生成角色提示词** — 为每个角色生成 4 个角度（正面/背面/左侧面/右侧面）
3. **生成场景提示词** — 描述场景的稳定图像生成提示词
4. **生成语音（TTS）** — 将对话台词转为语音描述
5. **导出 ComfyUI 工作流** — 生成 workflow JSON
## 规则
1. 先用中文与用户讨论脚本构思、视觉风格、角色设定
2. 用户确认后（说"开始制作"或"生成"），再分步输出结构化内容
3. **分两步走：第一步输出角色+场景+TTS，第二步输出 WORKFLOW**（不要一次全出）
## 输出格式
直接输出普通文本，用 emoji 标记不同内容类型：
- 🎭 角色名 — 角色描述和相关提示词
- 🌄 场景描述 — 场景提示词
- 🔊 角色名: "台词" — TTS 语音
- 📦 ComfyUI 工作流 — 最终的 workflow JSON（放在 ```json 代码块内）
## 结构化标记格式
严格按照以下格式输出，系统会自动解析并生成图片/TTS：
## 工作流程
1. 用户提供脚本 → 你分析脚本结构
2. 针对每个角色 → 生成 4 角度提示词
3. 针对每个场景 → 生成场景提示词
4. 针对每段对话 → 生成 TTS 语音描述
5. 最终导出 ComfyUI 工作流 JSON
【CHARACTER】
name: 林澈
description: 女宇航员，30岁，冷静坚毅，短发收在头盔内
prompt_front: cinematic full body front view of a female astronaut, white EVA suit, photorealistic
prompt_back: cinematic full body back view of a female astronaut, life support backpack
prompt_left: cinematic full body left side view of a female astronaut, side utility pouches
prompt_right: cinematic full body right side view of a female astronaut, helmet visor reflection
【/CHARACTER】
始终用中文回复，文字简洁清晰。"#;
【SCENE】
id: 1
description: 飞船驾驶舱内景
prompt: interior of a sci-fi spacecraft cockpit, holographic displays, cold blue lighting, cinematic
【/SCENE】
【TTS】
character: 林澈
text: 我们准备出舱检查
emotion: 平静
【/TTS】
【WORKFLOW】
```json
{"comfyui_workflow": {}}
```
【/WORKFLOW】
## 关键要求
- prompt_* 用英文，同时包含主体 + 视角 + 细节 + 画质标签
- 不要一次输出太多，先输角色和场景，下次再输工作流
- 如果没有用户明确要求，切勿自作主张输出全部内容
"#;
fn upstream_base() -> String {
std::env::var("BOOSAPI_UPSTREAM_BASE")
.unwrap_or_else(|_| "https://api.boos.lat/v1".to_string())
.trim_end_matches('/')
.to_string()
}
fn upstream_key() -> String {
std::env::var("BOOSAPI_UPSTREAM_KEY").unwrap_or_else(|_| "sk-boos4721".to_string())
}
pub fn routes() -> Router<AppState> {
Router::new().route("/api/comfyui-agent/chat", post(chat_handler))
}
#[derive(Debug, Default)]
struct ParsedSections {
characters: Vec<CharacterSection>,
scenes: Vec<SceneSection>,
tts_list: Vec<TtsSection>,
workflow_json: Option<String>,
text_content: String,
}
#[derive(Debug, Default)]
struct CharacterSection {
name: String,
description: String,
prompt_front: String,
prompt_back: String,
prompt_left: String,
prompt_right: String,
}
#[derive(Debug)]
struct SceneSection {
id: u32,
description: String,
prompt: String,
}
#[derive(Debug)]
struct TtsSection {
character: String,
text: String,
emotion: String,
}
fn parse_sections(full_text: &str) -> ParsedSections {
let mut sections = ParsedSections::default();
sections.text_content = full_text.to_string();
// Parse CHARACTER blocks
let char_re = Regex::new(r"(?s)【CHARACTER】\s*(.*?)【/CHARACTER】").unwrap();
for cap in char_re.captures_iter(full_text) {
let block = cap.get(1).unwrap().as_str();
let mut char_section = CharacterSection::default();
for line in block.lines() {
let line = line.trim();
let line_lower = line.to_lowercase();
if let Some(val) = line.strip_prefix("name:").or_else(|| line.strip_prefix("name :")) {
char_section.name = val.trim().to_string();
} else if let Some(val) = line.strip_prefix("description:").or_else(|| line.strip_prefix("description :")) {
char_section.description = val.trim().to_string();
} else if line_lower.starts_with("prompt_front:") || line_lower.starts_with("prompt_front :") {
let val = line.splitn(2, ':').nth(1).unwrap_or("").trim();
char_section.prompt_front = val.to_string();
} else if line_lower.starts_with("prompt_back:") || line_lower.starts_with("prompt_back :") {
let val = line.splitn(2, ':').nth(1).unwrap_or("").trim();
char_section.prompt_back = val.to_string();
} else if line_lower.starts_with("prompt_left:") || line_lower.starts_with("prompt_left :") {
let val = line.splitn(2, ':').nth(1).unwrap_or("").trim();
char_section.prompt_left = val.to_string();
} else if line_lower.starts_with("prompt_right:") || line_lower.starts_with("prompt_right :") {
let val = line.splitn(2, ':').nth(1).unwrap_or("").trim();
char_section.prompt_right = val.to_string();
}
}
if !char_section.name.is_empty() {
sections.characters.push(char_section);
}
}
// Parse SCENE blocks
let scene_re = Regex::new(r"(?s)【SCENE】\s*(.*?)【/SCENE】").unwrap();
for cap in scene_re.captures_iter(full_text) {
let block = cap.get(1).unwrap().as_str();
let mut scene = SceneSection {
id: 0,
description: String::new(),
prompt: String::new(),
};
for line in block.lines() {
let line = line.trim();
let line_lower = line.to_lowercase();
if line_lower.starts_with("id:") || line_lower.starts_with("id :") {
let val = line.splitn(2, ':').nth(1).unwrap_or("").trim();
scene.id = val.parse().unwrap_or(0);
} else if let Some(val) = line.strip_prefix("description:").or_else(|| line.strip_prefix("description :")) {
scene.description = val.trim().to_string();
} else if line_lower.starts_with("prompt:") || line_lower.starts_with("prompt :") {
let val = line.splitn(2, ':').nth(1).unwrap_or("").trim();
scene.prompt = val.to_string();
}
}
if !scene.prompt.is_empty() {
sections.scenes.push(scene);
}
}
// Parse TTS blocks
let tts_re = Regex::new(r"(?s)【TTS】\s*(.*?)【/TTS】").unwrap();
for cap in tts_re.captures_iter(full_text) {
let block = cap.get(1).unwrap().as_str();
let mut tts = TtsSection {
character: String::new(),
text: String::new(),
emotion: String::new(),
};
for line in block.lines() {
let line = line.trim();
let line_lower = line.to_lowercase();
if line_lower.starts_with("character:") || line_lower.starts_with("character :") {
let val = line.splitn(2, ':').nth(1).unwrap_or("").trim();
tts.character = val.to_string();
} else if line_lower.starts_with("text:") || line_lower.starts_with("text :") {
let val = line.splitn(2, ':').nth(1).unwrap_or("").trim();
tts.text = val.to_string();
} else if line_lower.starts_with("emotion:") || line_lower.starts_with("emotion :") {
let val = line.splitn(2, ':').nth(1).unwrap_or("").trim();
tts.emotion = val.to_string();
}
}
if !tts.text.is_empty() {
sections.tts_list.push(tts);
}
}
// Parse WORKFLOW blocks
let wf_re = Regex::new(r"(?s)【WORKFLOW】\s*```(?:json)?\s*(.*?)```\s*【/WORKFLOW】").unwrap();
if let Some(cap) = wf_re.captures(full_text) {
let wf = cap.get(1).unwrap().as_str().trim().to_string();
sections.workflow_json = Some(wf);
}
sections
}
async fn call_image_gen(prompt: &str) -> Result<String, String> {
let client = Client::new();
let url = format!("{}/images/generations", upstream_base());
let resp = client
.post(&url)
.bearer_auth(upstream_key())
.json(&json!({
"model": "gpt-image-2",
"prompt": prompt,
"n": 1,
"size": "1024x1024"
}))
.send()
.await
.map_err(|e| format!("image request failed: {}", e))?;
let status = resp.status();
let data: Value = resp
.json()
.await
.map_err(|e| format!("image response parse failed: {}", e))?;
if !status.is_success() {
return Err(format!("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(text: &str, voice: &str) -> Result<Vec<u8>, String> {
let client = Client::new();
let url = format!("{}/audio/speech", upstream_base());
let resp = client
.post(&url)
.bearer_auth(upstream_key())
.json(&json!({
"model": "tts-1",
"input": text,
"voice": voice
}))
.send()
.await
.map_err(|e| format!("TTS request failed: {}", e))?;
let status = resp.status();
if !status.is_success() {
let body = resp.text().await.unwrap_or_default();
return Err(format!("TTS upstream returned {}: {}", status, body));
}
let bytes = resp
.bytes()
.await
.map_err(|e| format!("TTS response read failed: {}", e))?;
Ok(bytes.to_vec())
}
fn send_sse(tx: &UnboundedSender<Result<Event, Infallible>>, payload: Value) -> bool {
tx.send(Ok(Event::default().data(payload.to_string()))).is_ok()
}
async fn chat_handler(
State(state): State<AppState>,
headers: HeaderMap,
@@ -91,12 +325,7 @@ async fn chat_handler(
let (tx, rx) = unbounded_channel::<Result<Event, Infallible>>();
tokio::spawn(async move {
let upstream_base = std::env::var("BOOSAPI_UPSTREAM_BASE")
.unwrap_or_else(|_| "https://api.boos.lat/v1".to_string());
let upstream_key =
std::env::var("BOOSAPI_UPSTREAM_KEY").unwrap_or_else(|_| "sk-boos4721".to_string());
let upstream_url = format!("{}/chat/completions", upstream_base.trim_end_matches('/'));
let upstream_url = format!("{}/chat/completions", upstream_base().trim_end_matches('/'));
let mut upstream_messages = Vec::with_capacity(messages.len() + 1);
upstream_messages.push(json!({
@@ -107,11 +336,11 @@ async fn chat_handler(
let client = Client::new();
info!("comfyui-agent chat streaming request to upstream");
info!("comfyui-agent: sending chat request to upstream");
let response = match client
.post(upstream_url)
.bearer_auth(upstream_key)
.post(&upstream_url)
.bearer_auth(upstream_key())
.json(&json!({
"model": "gpt-5.5",
"messages": upstream_messages,
@@ -122,12 +351,12 @@ async fn chat_handler(
.await
{
Ok(response) => response,
Err(_) => {
Err(e) => {
send_sse(
&tx,
json!({
"type": "text",
"content": "服务暂时不可用，请稍后重试",
"content": format!("服务暂时不可用，请稍后重试。({})", e),
}),
);
send_sse(&tx, json!({ "type": "done" }));
@@ -149,7 +378,9 @@ async fn chat_handler(
let mut stream = response.bytes_stream();
let mut buffer: Vec<u8> = Vec::new();
let mut accumulated_text = String::new();
// Phase 1: Stream the LLM text response
loop {
let chunk = tokio::select! {
_ = tx.closed() => {
@@ -165,15 +396,13 @@ async fn chat_handler(
let chunk = match chunk {
Ok(chunk) => chunk,
Err(error) => {
if !send_sse(
send_sse(
&tx,
json!({
"type": "text",
"content": format!("\n\n[错误: {}]", error),
}),
) {
return;
}
);
break;
}
};
@@ -191,19 +420,19 @@ async fn chat_handler(
}
let line = String::from_utf8_lossy(&line_bytes);
let line = line.trim();
let line = line.trim().to_string();
if !line.starts_with("data: ") {
continue;
}
let data = line.trim_start_matches("data: ").trim();
let data = line.trim_start_matches("data: ").trim().to_string();
if data == "[DONE]" {
continue;
}
let Ok(parsed) = serde_json::from_str::<Value>(data) else {
let Ok(parsed) = serde_json::from_str::<Value>(&data) else {
continue;
};
@@ -214,29 +443,216 @@ async fn chat_handler(
.and_then(|choice| choice.get("delta"))
.and_then(|delta| delta.get("content"))
.and_then(Value::as_str)
.unwrap_or_default();
.unwrap_or_default()
.to_string();
if !content.is_empty()
&& !send_sse(
if !content.is_empty() {
accumulated_text.push_str(&content);
if !send_sse(
&tx,
json!({
"type": "text",
"content": content,
}),
)
{
return;
) {
return;
}
}
}
}
// Phase 2: Parse structured sections from accumulated text
info!("comfyui-agent: text streaming complete, parsing sections");
let sections = parse_sections(&accumulated_text);
let has_structured_data = !sections.characters.is_empty()
|| !sections.scenes.is_empty()
|| !sections.tts_list.is_empty()
|| sections.workflow_json.is_some();
if !has_structured_data {
send_sse(&tx, json!({ "type": "done" }));
return;
}
// Send a processing indicator
send_sse(
&tx,
json!({
"type": "text",
"content": "\n\n🔄 正在处理结构化数据..."
}),
);
// Generate images for each character (4 angles each)
for (char_idx, character) in sections.characters.iter().enumerate() {
// Send character_prompts event
send_sse(
&tx,
json!({
"type": "character_prompts",
"character": character.name,
}),
);
let angles = [
("正面", &character.prompt_front, "front"),
("背面", &character.prompt_back, "back"),
("左侧", &character.prompt_left, "left"),
("右侧", &character.prompt_right, "right"),
];
for (angle_label, prompt, angle_key) in &angles {
if prompt.is_empty() {
continue;
}
info!(
"comfyui-agent: generating image for {} - {}",
character.name, angle_label
);
// Add a short delay between image gen requests to avoid rate limiting
if char_idx > 0 || *angle_key != "front" {
tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
}
match call_image_gen(prompt).await {
Ok(url) if !url.is_empty() => {
send_sse(
&tx,
json!({
"type": "image",
"url": url,
"character": character.name,
"angle": angle_key,
}),
);
}
_ => {
send_sse(
&tx,
json!({
"type": "text",
"content": format!("\n⚠️ {} [{}] 图片生成失败，请重试", character.name, angle_label),
}),
);
}
}
}
}
// Generate images for each scene
for scene in &sections.scenes {
info!("comfyui-agent: generating scene image for scene {}", scene.id);
tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
match call_image_gen(&scene.prompt).await {
Ok(url) if !url.is_empty() => {
send_sse(
&tx,
json!({
"type": "scene",
"id": scene.id,
"description": scene.description,
"prompt": scene.prompt,
"url": url,
}),
);
}
_ => {
// Still send scene info even without image
send_sse(
&tx,
json!({
"type": "scene",
"id": scene.id,
"description": scene.description,
"prompt": scene.prompt,
}),
);
}
}
}
// Generate TTS for each dialogue
for tts in &sections.tts_list {
info!("comfyui-agent: generating TTS for {}", tts.character);
// Choose voice based on emotion
let voice = match tts.emotion.as_str() {
"激动" | "愤怒" => "nova",
"悲伤" | "恐惧" => "shimmer",
_ => "alloy",
};
match call_tts(&tts.text, voice).await {
Ok(audio_bytes) => {
// For SSE, we'll send a data URL for small audio, or a placeholder
// Since SSE can't easily send binary, we store to file and send URL
let audio_path = format!("/tmp/comfyui_tts_{}_{}.mp3",
tts.character.replace(' ', "_"),
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_millis()
);
if let Ok(()) = tokio::fs::write(&audio_path, &audio_bytes).await {
send_sse(
&tx,
json!({
"type": "tts",
"character": tts.character,
"text": tts.text,
"url": format!("/audio/{}", audio_path.trim_start_matches("/tmp/comfyui_tts_")),
}),
);
}
}
Err(e) => {
info!("comfyui-agent: TTS generation failed: {}", e);
send_sse(
&tx,
json!({
"type": "text",
"content": format!("\n⚠️ {} 的语音生成失败", tts.character),
}),
);
}
}
tokio::time::sleep(tokio::time::Duration::from_millis(500)).await;
}
// Send ComfyUI workflow if present
if let Some(wf_json) = &sections.workflow_json {
info!("comfyui-agent: sending workflow JSON");
send_sse(
&tx,
json!({
"type": "comfyui_workflow",
"content": wf_json,
}),
);
}
// Signal done with summary
let mut summary = String::from("\n\n✅ **处理完成！**");
if !sections.characters.is_empty() {
summary.push_str(&format!("\n- 角色: {} 个", sections.characters.len()));
}
if !sections.scenes.is_empty() {
summary.push_str(&format!("\n- 场景: {} 个", sections.scenes.len()));
}
if !sections.tts_list.is_empty() {
summary.push_str(&format!("\n- TTS: {} 段", sections.tts_list.len()));
}
if sections.workflow_json.is_some() {
summary.push_str("\n- ComfyUI 工作流: ✅ 已生成");
}
send_sse(&tx, json!({ "type": "text", "content": summary }));
send_sse(&tx, json!({ "type": "done" }));
});
Ok(Sse::new(UnboundedReceiverStream::new(rx)))
}
fn send_sse(tx: &UnboundedSender<Result<Event, Infallible>>, payload: Value) -> bool {
tx.send(Ok(Event::default().data(payload.to_string())))
.is_ok()
}
boosapi-rs/src/routes/mod.rs
+51 -3
View File
@@ -1,5 +1,14 @@
use crate::AppState;
use axum::{routing::get, Router};
use axum::{
body::Body,
http::{Request, StatusCode, Uri},
response::{Html, Response},
routing::get,
Router,
};
use std::convert::Infallible;
use tower::service_fn;
use tower_http::services::ServeDir;
mod admin;
mod auth;
@@ -22,7 +31,11 @@ mod user_api_keys;
mod video_agent;
pub fn build_router(state: AppState) -> Router {
Router::new()
let web_dir = std::env::var("BOOSAPI_WEB_DIR")
.unwrap_or_else(|_| "../dist/web".to_string());
// Build the API router
let api_routes = Router::new()
.route("/healthz", get(health::health_check))
.merge(admin::routes(state.clone()))
.merge(proxy::routes())
@@ -42,5 +55,40 @@ pub fn build_router(state: AppState) -> Router {
.merge(user::routes())
.merge(user_api_keys::routes())
.merge(video_agent::routes())
.with_state(state)
.with_state(state);
// Static file serving with SPA fallback
let web_dir_owned = web_dir.clone();
let static_service = ServeDir::new(&web_dir)
.append_index_html_on_directories(true)
.fallback(service_fn(move |req: Request<Body>| {
let web_dir = web_dir_owned.clone();
async move {
let index_path = format!("{}/index.html", web_dir);
match tokio::fs::read_to_string(&index_path).await {
Ok(html) => {
Ok::<_, Infallible>(
Response::builder()
.status(StatusCode::OK)
.header("Content-Type", "text/html; charset=utf-8")
.body(Body::from(html))
.unwrap(),
)
}
Err(_) => {
Ok::<_, Infallible>(
Response::builder()
.status(StatusCode::NOT_FOUND)
.header("Content-Type", "text/plain; charset=utf-8")
.body(Body::from(format!("404: {} not found", req.uri().path())))
.unwrap(),
)
}
}
}
}));
Router::new()
.merge(api_routes)
.fallback_service(static_service)
}
boosapi-rs/src/routes/proxy/mod.rs
+1 -1
View File
@@ -249,7 +249,7 @@ async fn list_models(State(state): State<AppState>, headers: HeaderMap) -> AppRe
check_auth(&state, &headers).await?;
let rows: Vec<(String, String)> = sqlx::query_as(
"SELECT tr.name, tr.model FROM token_routes tr WHERE tr.enabled = 1 ORDER BY tr.name",
"SELECT tr.model_pattern, COALESCE(tr.model_mapping, tr.model_pattern) FROM token_routes tr WHERE tr.enabled = 1 ORDER BY tr.id",
)
.fetch_all(&*state.db)
.await?;
boosapi-rs/src/routes/user/mod.rs
+1 -1
View File
@@ -7,7 +7,7 @@ use crate::{
use axum::{
extract::{Path, State},
http::{HeaderMap, StatusCode},
routing::{delete, get, post, put},
routing::{delete, get},
Json, Router,
};
use serde::Deserialize;
boosapi-rs/src/services/billing/mod.rs
+282
View File
@@ -0,0 +1,282 @@
use crate::{
error::{AppError, AppResult},
services::shared::{enforce_key_request_limit, increment_key_usage},
};
use serde::{Deserialize, Serialize};
use sqlx::{FromRow, SqlitePool};
use std::sync::Arc;
#[derive(Clone)]
pub struct UsageTracker {
db: Arc<SqlitePool>,
}
#[derive(Debug, Clone, Serialize, Deserialize, FromRow, PartialEq)]
pub struct UsageSummary {
pub downstream_key_id: i64,
pub request_count: i64,
pub prompt_tokens: i64,
pub completion_tokens: i64,
pub total_tokens: i64,
pub total_cost: f64,
}
impl UsageTracker {
pub fn new(db: Arc<SqlitePool>) -> Self {
Self { db }
}
pub async fn record_usage(
&self,
key_id: i64,
model: impl Into<String>,
prompt_tokens: i64,
completion_tokens: i64,
cost: f64,
) -> AppResult<()> {
Self::ensure_schema(&self.db).await?;
let prompt_tokens = prompt_tokens.max(0);
let completion_tokens = completion_tokens.max(0);
let total_tokens = prompt_tokens + completion_tokens;
sqlx::query(
"INSERT INTO usage_logs
(downstream_key_id, model, prompt_tokens, completion_tokens, total_tokens, cost)
VALUES (?, ?, ?, ?, ?, ?)",
)
.bind(key_id)
.bind(model.into())
.bind(prompt_tokens)
.bind(completion_tokens)
.bind(total_tokens)
.bind(cost.max(0.0))
.execute(&*self.db)
.await?;
increment_key_usage(&self.db, key_id, cost.max(0.0)).await?;
Ok(())
}
pub async fn get_key_usage(
&self,
key_id: i64,
since: Option<&str>,
until: Option<&str>,
) -> AppResult<UsageSummary> {
Self::ensure_schema(&self.db).await?;
let summary = sqlx::query_as::<_, UsageSummary>(
"SELECT
? AS downstream_key_id,
COUNT(*) AS request_count,
COALESCE(SUM(prompt_tokens), 0) AS prompt_tokens,
COALESCE(SUM(completion_tokens), 0) AS completion_tokens,
COALESCE(SUM(total_tokens), 0) AS total_tokens,
COALESCE(SUM(cost), 0.0) AS total_cost
FROM usage_logs
WHERE downstream_key_id = ?
AND (? IS NULL OR created_at >= ?)
AND (? IS NULL OR created_at <= ?)",
)
.bind(key_id)
.bind(key_id)
.bind(since)
.bind(since)
.bind(until)
.bind(until)
.fetch_one(&*self.db)
.await?;
Ok(summary)
}
pub async fn check_rate_limit(&self, key_id: i64) -> AppResult<()> {
enforce_key_request_limit(&self.db, key_id, 1).await?;
let row: Option<(Option<f64>, Option<f64>, Option<bool>)> = sqlx::query_as(
"SELECT max_cost, used_cost, enabled
FROM downstream_api_keys
WHERE id = ?",
)
.bind(key_id)
.fetch_optional(&*self.db)
.await?;
let (max_cost, used_cost, enabled) =
row.ok_or_else(|| AppError::Unauthorized("API key not found".into()))?;
if !enabled.unwrap_or(true) {
return Err(AppError::Forbidden("API key is disabled".into()));
}
if let Some(max_cost) = max_cost {
if used_cost.unwrap_or(0.0) >= max_cost {
return Err(AppError::RateLimited);
}
}
Ok(())
}
pub async fn create_daily_usage_snapshot(&self) -> AppResult<()> {
Self::create_daily_usage_snapshot_for_db(&self.db).await
}
pub async fn create_daily_usage_snapshot_for_db(db: &SqlitePool) -> AppResult<()> {
Self::ensure_schema(db).await?;
sqlx::query(
"INSERT INTO usage_daily_snapshots
(local_day, downstream_key_id, request_count, prompt_tokens,
completion_tokens, total_tokens, total_cost, updated_at)
SELECT
date(created_at, 'localtime') AS local_day,
downstream_key_id,
COUNT(*) AS request_count,
COALESCE(SUM(prompt_tokens), 0) AS prompt_tokens,
COALESCE(SUM(completion_tokens), 0) AS completion_tokens,
COALESCE(SUM(total_tokens), 0) AS total_tokens,
COALESCE(SUM(cost), 0.0) AS total_cost,
datetime('now') AS updated_at
FROM usage_logs
GROUP BY date(created_at, 'localtime'), downstream_key_id
ON CONFLICT(local_day, downstream_key_id) DO UPDATE SET
request_count = excluded.request_count,
prompt_tokens = excluded.prompt_tokens,
completion_tokens = excluded.completion_tokens,
total_tokens = excluded.total_tokens,
total_cost = excluded.total_cost,
updated_at = datetime('now')",
)
.execute(db)
.await?;
Ok(())
}
async fn ensure_schema(db: &SqlitePool) -> AppResult<()> {
sqlx::query(
"CREATE TABLE IF NOT EXISTS usage_logs (
id INTEGER PRIMARY KEY AUTOINCREMENT,
downstream_key_id INTEGER,
model TEXT,
prompt_tokens INTEGER DEFAULT 0,
completion_tokens INTEGER DEFAULT 0,
total_tokens INTEGER DEFAULT 0,
cost REAL DEFAULT 0.0,
created_at TEXT DEFAULT (datetime('now'))
)",
)
.execute(db)
.await?;
sqlx::query(
"CREATE INDEX IF NOT EXISTS usage_logs_downstream_key_created_idx
ON usage_logs (downstream_key_id, created_at)",
)
.execute(db)
.await?;
sqlx::query(
"CREATE TABLE IF NOT EXISTS usage_daily_snapshots (
id INTEGER PRIMARY KEY AUTOINCREMENT,
local_day TEXT NOT NULL,
downstream_key_id INTEGER,
request_count INTEGER NOT NULL DEFAULT 0,
prompt_tokens INTEGER NOT NULL DEFAULT 0,
completion_tokens INTEGER NOT NULL DEFAULT 0,
total_tokens INTEGER NOT NULL DEFAULT 0,
total_cost REAL NOT NULL DEFAULT 0.0,
created_at TEXT DEFAULT (datetime('now')),
updated_at TEXT DEFAULT (datetime('now')),
UNIQUE(local_day, downstream_key_id)
)",
)
.execute(db)
.await?;
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
use sqlx::sqlite::SqlitePoolOptions;
async fn pool() -> SqlitePool {
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,
enabled INTEGER DEFAULT 1,
max_cost REAL,
used_cost REAL DEFAULT 0,
max_requests INTEGER,
used_requests INTEGER DEFAULT 0,
last_used_at TEXT,
updated_at TEXT DEFAULT (datetime('now'))
)",
)
.execute(&pool)
.await
.unwrap();
sqlx::query(
"INSERT INTO downstream_api_keys
(id, name, key, enabled, max_cost, used_cost, max_requests, used_requests)
VALUES (1, 'test', 'sk-test', 1, 1.0, 0.0, 5, 0)",
)
.execute(&pool)
.await
.unwrap();
pool
}
#[tokio::test]
async fn records_and_aggregates_usage() {
let tracker = UsageTracker::new(Arc::new(pool().await));
tracker
.record_usage(1, "gpt-test", 10, 20, 0.25)
.await
.unwrap();
let summary = tracker.get_key_usage(1, None, None).await.unwrap();
assert_eq!(summary.request_count, 1);
assert_eq!(summary.prompt_tokens, 10);
assert_eq!(summary.completion_tokens, 20);
assert_eq!(summary.total_tokens, 30);
assert_eq!(summary.total_cost, 0.25);
}
#[tokio::test]
async fn rate_limit_checks_request_and_cost_limits() {
let tracker = UsageTracker::new(Arc::new(pool().await));
tracker.check_rate_limit(1).await.unwrap();
sqlx::query("UPDATE downstream_api_keys SET used_requests = 5 WHERE id = 1")
.execute(&*tracker.db)
.await
.unwrap();
assert!(matches!(
tracker.check_rate_limit(1).await,
Err(AppError::RateLimited)
));
}
#[tokio::test]
async fn creates_daily_snapshot() {
let tracker = UsageTracker::new(Arc::new(pool().await));
tracker
.record_usage(1, "gpt-test", 2, 3, 0.1)
.await
.unwrap();
tracker.create_daily_usage_snapshot().await.unwrap();
let count: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM usage_daily_snapshots")
.fetch_one(&*tracker.db)
.await
.unwrap();
assert_eq!(count, 1);
}
}
boosapi-rs/src/services/cron/mod.rs
+2 -1
View File
@@ -1,4 +1,4 @@
use crate::{error::AppResult, proxy::token_router::TokenRouter};
use crate::{error::AppResult, proxy::token_router::TokenRouter, services::billing::UsageTracker};
use sqlx::SqlitePool;
use std::sync::{
atomic::{AtomicBool, Ordering},
@@ -78,6 +78,7 @@ impl CronScheduler {
}
async fn aggregate_usage(db: Arc<SqlitePool>) -> AppResult<()> {
UsageTracker::create_daily_usage_snapshot_for_db(&db).await?;
sqlx::query(
"UPDATE downstream_api_keys
SET updated_at = datetime('now')
boosapi-rs/src/services/mod.rs
+2
View File
@@ -1,5 +1,7 @@
pub mod announcement;
pub mod billing;
pub mod cron;
pub mod notification;
pub mod oauth;
pub mod shared;
pub mod update_center;
boosapi-rs/src/services/notification/mod.rs
+544
View File
@@ -0,0 +1,544 @@
use crate::error::AppResult;
use chrono::{FixedOffset, SecondsFormat, TimeZone, Utc};
use lettre::{
message::Mailbox, transport::smtp::authentication::Credentials, Message, SmtpTransport,
Transport,
};
use reqwest::Url;
use serde::{Deserialize, Serialize};
use serde_json::json;
use sqlx::SqlitePool;
use std::{collections::HashMap, sync::Arc};
#[derive(Clone)]
pub struct NotificationService {
pub db: Arc<SqlitePool>,
pub client: reqwest::Client,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "lowercase")]
pub enum NotificationLevel {
Info,
Warning,
Error,
Critical,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Notification {
pub title: String,
pub message: String,
pub level: NotificationLevel,
pub channel: Option<String>,
pub signature: Option<String>,
}
#[derive(Debug, Clone, Default)]
struct NotificationConfig {
telegram_bot_token: Option<String>,
telegram_chat_id: Option<String>,
telegram_thread_id: Option<i64>,
wecom_webhook: Option<String>,
feishu_webhook: Option<String>,
webhook_url: Option<String>,
bark_url: Option<String>,
serverchan_key: Option<String>,
smtp: SmtpConfig,
cooldown_sec: i64,
}
#[derive(Debug, Clone, Default)]
struct SmtpConfig {
host: Option<String>,
port: Option<u16>,
user: Option<String>,
pass: Option<String>,
from: Option<String>,
to: Option<String>,
}
impl NotificationService {
pub fn new(db: Arc<SqlitePool>) -> Self {
Self {
db,
client: reqwest::Client::new(),
}
}
pub fn with_client(db: Arc<SqlitePool>, client: reqwest::Client) -> Self {
Self { db, client }
}
pub async fn send(&self, notif: Notification) -> AppResult<()> {
Self::ensure_schema(&self.db).await?;
let config = self.load_config().await?;
if self.should_throttle(&notif, config.cooldown_sec).await? {
return Ok(());
}
let mut channels = self.enabled_channels(&config);
if let Some(channel) = notif.channel.as_deref() {
channels.retain(|candidate| candidate == &normalize_channel(channel));
}
if channels.is_empty() {
tracing::debug!(
channel = notif.channel.as_deref().unwrap_or("all"),
"notification skipped because no matching channel is configured"
);
return Ok(());
}
for channel in channels {
if let Err(err) = self.send_channel(&channel, &notif, &config).await {
tracing::warn!(%err, channel, "notification channel failed");
}
}
Ok(())
}
async fn load_config(&self) -> AppResult<NotificationConfig> {
let rows: Vec<(String, Option<String>)> =
sqlx::query_as("SELECT key, value FROM settings WHERE key LIKE 'notify_%'")
.fetch_all(&*self.db)
.await?;
let map: HashMap<String, String> = rows
.into_iter()
.filter_map(|(key, value)| value.map(|value| (key, value)))
.filter(|(_, value)| !value.trim().is_empty())
.collect();
Ok(NotificationConfig {
telegram_bot_token: setting(&map, "notify_telegram_bot_token"),
telegram_chat_id: setting(&map, "notify_telegram_chat_id"),
telegram_thread_id: setting(&map, "notify_telegram_thread_id")
.and_then(|value| value.parse().ok()),
wecom_webhook: setting(&map, "notify_wecom_webhook"),
feishu_webhook: setting(&map, "notify_feishu_webhook"),
webhook_url: setting(&map, "notify_webhook_url"),
bark_url: setting(&map, "notify_bark_url"),
serverchan_key: setting(&map, "notify_serverchan_key"),
smtp: SmtpConfig {
host: setting(&map, "notify_smtp_host"),
port: setting(&map, "notify_smtp_port").and_then(|value| value.parse().ok()),
user: setting(&map, "notify_smtp_user"),
pass: setting(&map, "notify_smtp_pass"),
from: setting(&map, "notify_smtp_from"),
to: setting(&map, "notify_smtp_to"),
},
cooldown_sec: setting(&map, "notify_cooldown_sec")
.and_then(|value| value.parse().ok())
.unwrap_or(60),
})
}
async fn should_throttle(&self, notif: &Notification, cooldown_sec: i64) -> AppResult<bool> {
let Some(signature) = notif.signature.as_deref().map(str::trim) else {
return Ok(false);
};
if signature.is_empty() || cooldown_sec <= 0 {
return Ok(false);
}
let updated = sqlx::query(
"UPDATE notification_throttle
SET count = count + 1,
title = ?,
message = ?,
level = ?,
updated_at = datetime('now')
WHERE signature = ?
AND updated_at >= datetime('now', ?)",
)
.bind(&notif.title)
.bind(&notif.message)
.bind(level_name(&notif.level))
.bind(signature)
.bind(format!("-{} seconds", cooldown_sec))
.execute(&*self.db)
.await?;
if updated.rows_affected() > 0 {
tracing::info!(signature, "notification merged into cooldown throttle");
return Ok(true);
}
sqlx::query(
"INSERT INTO notification_throttle
(signature, title, message, level, count, updated_at)
VALUES (?, ?, ?, ?, 1, datetime('now'))
ON CONFLICT(signature) DO UPDATE SET
title = excluded.title,
message = excluded.message,
level = excluded.level,
count = 1,
updated_at = datetime('now')",
)
.bind(signature)
.bind(&notif.title)
.bind(&notif.message)
.bind(level_name(&notif.level))
.execute(&*self.db)
.await?;
Ok(false)
}
fn enabled_channels(&self, config: &NotificationConfig) -> Vec<String> {
let mut channels = Vec::new();
if config.telegram_bot_token.is_some() && config.telegram_chat_id.is_some() {
channels.push("telegram".to_string());
}
if config.wecom_webhook.is_some() {
channels.push("wecom".to_string());
}
if config.feishu_webhook.is_some() {
channels.push("feishu".to_string());
}
if config.webhook_url.is_some() {
channels.push("webhook".to_string());
}
if config.bark_url.is_some() {
channels.push("bark".to_string());
}
if config.serverchan_key.is_some() {
channels.push("serverchan".to_string());
}
if config.smtp.is_configured() {
channels.push("smtp".to_string());
}
channels
}
async fn send_channel(
&self,
channel: &str,
notif: &Notification,
config: &NotificationConfig,
) -> anyhow::Result<()> {
match channel {
"telegram" => self.send_telegram(notif, config).await,
"wecom" => self.send_wecom(notif, config).await,
"feishu" => self.send_feishu(notif, config).await,
"webhook" => self.send_webhook(notif, config).await,
"bark" => self.send_bark(notif, config).await,
"serverchan" => self.send_serverchan(notif, config).await,
"smtp" => self.send_smtp(notif, config).await,
_ => Ok(()),
}
}
async fn send_telegram(
&self,
notif: &Notification,
config: &NotificationConfig,
) -> anyhow::Result<()> {
let token = required(config.telegram_bot_token.as_deref(), "telegram token")?;
let chat_id = required(config.telegram_chat_id.as_deref(), "telegram chat id")?;
let mut body = json!({
"chat_id": chat_id,
"text": truncate_chars(&format_notification_text(notif), 3900),
});
if let Some(thread_id) = config.telegram_thread_id {
body["message_thread_id"] = json!(thread_id);
}
self.client
.post(format!("https://api.telegram.org/bot{token}/sendMessage"))
.json(&body)
.send()
.await?
.error_for_status()?;
Ok(())
}
async fn send_wecom(
&self,
notif: &Notification,
config: &NotificationConfig,
) -> anyhow::Result<()> {
let url = required(config.wecom_webhook.as_deref(), "wecom webhook")?;
self.client
.post(url)
.json(&json!({
"msgtype": "text",
"text": { "content": format_notification_text(notif) },
}))
.send()
.await?
.error_for_status()?;
Ok(())
}
async fn send_feishu(
&self,
notif: &Notification,
config: &NotificationConfig,
) -> anyhow::Result<()> {
let url = required(config.feishu_webhook.as_deref(), "feishu webhook")?;
self.client
.post(url)
.json(&json!({
"msg_type": "text",
"content": { "text": format_notification_text(notif) },
}))
.send()
.await?
.error_for_status()?;
Ok(())
}
async fn send_webhook(
&self,
notif: &Notification,
config: &NotificationConfig,
) -> anyhow::Result<()> {
let url = required(config.webhook_url.as_deref(), "webhook url")?;
self.client
.post(url)
.json(&json!({
"title": &notif.title,
"message": &notif.message,
"level": level_name(&notif.level),
}))
.send()
.await?
.error_for_status()?;
Ok(())
}
async fn send_bark(
&self,
notif: &Notification,
config: &NotificationConfig,
) -> anyhow::Result<()> {
let base = required(config.bark_url.as_deref(), "bark url")?;
let mut url = Url::parse(base)?;
{
let mut segments = url
.path_segments_mut()
.map_err(|_| anyhow::anyhow!("bark url cannot be a base url"))?;
segments.pop_if_empty();
segments.push(&notif.title);
segments.push(&notif.message);
}
self.client.get(url).send().await?.error_for_status()?;
Ok(())
}
async fn send_serverchan(
&self,
notif: &Notification,
config: &NotificationConfig,
) -> anyhow::Result<()> {
let key = required(config.serverchan_key.as_deref(), "serverchan key")?;
self.client
.post(format!("https://sctapi.ftqq.com/{key}.send"))
.form(&[
("title", notif.title.as_str()),
("desp", notif.message.as_str()),
])
.send()
.await?
.error_for_status()?;
Ok(())
}
async fn send_smtp(
&self,
notif: &Notification,
config: &NotificationConfig,
) -> anyhow::Result<()> {
let smtp = config.smtp.clone();
let notif = notif.clone();
tokio::task::spawn_blocking(move || smtp.send(&notif)).await??;
Ok(())
}
async fn ensure_schema(db: &SqlitePool) -> AppResult<()> {
sqlx::query(
"CREATE TABLE IF NOT EXISTS notification_throttle (
signature TEXT PRIMARY KEY NOT NULL,
title TEXT NOT NULL,
message TEXT NOT NULL,
level TEXT NOT NULL,
count INTEGER NOT NULL DEFAULT 1,
updated_at TEXT DEFAULT (datetime('now'))
)",
)
.execute(db)
.await?;
Ok(())
}
}
impl SmtpConfig {
fn is_configured(&self) -> bool {
self.host.is_some()
&& self.port.is_some()
&& self.user.is_some()
&& self.pass.is_some()
&& self.from.is_some()
}
fn send(self, notif: &Notification) -> anyhow::Result<()> {
let host = required(self.host.as_deref(), "smtp host")?;
let port = self.port.unwrap_or(587);
let user = required(self.user.as_deref(), "smtp user")?;
let pass = required(self.pass.as_deref(), "smtp pass")?;
let from = required(self.from.as_deref(), "smtp from")?;
let to = self.to.as_deref().unwrap_or(from);
let email = Message::builder()
.from(from.parse::<Mailbox>()?)
.to(to.parse::<Mailbox>()?)
.subject(format!(
"[BoosAPI][{}] {}",
level_name(&notif.level),
notif.title
))
.body(format_notification_text(notif))?;
let mailer = SmtpTransport::relay(host)?
.port(port)
.credentials(Credentials::new(user.to_string(), pass.to_string()))
.build();
mailer.send(&email)?;
Ok(())
}
}
fn setting(map: &HashMap<String, String>, key: &str) -> Option<String> {
map.get(key).map(|value| value.trim().to_string())
}
fn required<'a>(value: Option<&'a str>, name: &str) -> anyhow::Result<&'a str> {
value.ok_or_else(|| anyhow::anyhow!("{name} is not configured"))
}
fn normalize_channel(channel: &str) -> String {
channel.trim().to_ascii_lowercase()
}
fn level_name(level: &NotificationLevel) -> &'static str {
match level {
NotificationLevel::Info => "info",
NotificationLevel::Warning => "warning",
NotificationLevel::Error => "error",
NotificationLevel::Critical => "critical",
}
}
fn format_notification_text(notif: &Notification) -> String {
let utc = Utc::now();
let shanghai = FixedOffset::east_opt(8 * 3600)
.expect("valid Shanghai offset")
.from_utc_datetime(&utc.naive_utc());
format!(
"[BoosAPI][{}] {}\n\n{}\n\nLevel: {}\nLocal Time: {} (Asia/Shanghai)\nUTC Time: {}",
level_name(&notif.level).to_ascii_uppercase(),
notif.title,
notif.message,
level_name(&notif.level),
shanghai.format("%Y-%m-%d %H:%M:%S"),
utc.to_rfc3339_opts(SecondsFormat::Millis, true)
)
}
fn truncate_chars(value: &str, max_chars: usize) -> String {
if value.chars().count() <= max_chars {
return value.to_string();
}
value.chars().take(max_chars).collect()
}
#[cfg(test)]
mod tests {
use super::*;
use sqlx::sqlite::SqlitePoolOptions;
async fn pool() -> SqlitePool {
let pool = SqlitePoolOptions::new()
.max_connections(1)
.connect("sqlite::memory:")
.await
.unwrap();
sqlx::query("CREATE TABLE settings (key TEXT PRIMARY KEY NOT NULL, value TEXT)")
.execute(&pool)
.await
.unwrap();
pool
}
async fn set(pool: &SqlitePool, key: &str, value: &str) {
sqlx::query("INSERT OR REPLACE INTO settings (key, value) VALUES (?, ?)")
.bind(key)
.bind(value)
.execute(pool)
.await
.unwrap();
}
#[tokio::test]
async fn skips_when_no_channel_is_configured() {
let pool = pool().await;
let service = NotificationService::new(Arc::new(pool));
service
.send(Notification {
title: "Disk".to_string(),
message: "Almost full".to_string(),
level: NotificationLevel::Warning,
channel: Some("webhook".to_string()),
signature: None,
})
.await
.unwrap();
}
#[tokio::test]
async fn throttles_duplicate_signature() {
let pool = pool().await;
set(&pool, "notify_cooldown_sec", "60").await;
let service = NotificationService::new(Arc::new(pool));
let notification = Notification {
title: "Route down".to_string(),
message: "primary failed".to_string(),
level: NotificationLevel::Error,
channel: Some("webhook".to_string()),
signature: Some("route:1".to_string()),
};
service.send(notification.clone()).await.unwrap();
service.send(notification).await.unwrap();
let count: i64 = sqlx::query_scalar(
"SELECT count FROM notification_throttle WHERE signature = 'route:1'",
)
.fetch_one(&*service.db)
.await
.unwrap();
assert_eq!(count, 2);
}
#[test]
fn formats_telegram_compatible_text() {
let text = format_notification_text(&Notification {
title: "Title".to_string(),
message: "Message".to_string(),
level: NotificationLevel::Critical,
channel: None,
signature: None,
});
assert!(text.starts_with("[BoosAPI][CRITICAL] Title\n\nMessage"));
assert!(text.contains("Local Time:"));
assert!(text.contains("UTC Time:"));
}
#[test]
fn truncates_by_char_count() {
assert_eq!(truncate_chars("abcdef", 3), "abc");
assert_eq!(truncate_chars("短消息", 2), "短消");
}
}