feat(jimeng): add image/video generation routes with background jobs, workspace transfer and billing

This commit is contained in:
boos
2026-06-23 12:46:40 +08:00
parent 7087c75162
commit 7b68855c56
3 changed files with 688 additions and 1 deletions
+685
View File
@@ -0,0 +1,685 @@
//! Jimeng (Volcengine Visual / CV) async image & video generation routes.
//!
//! Mirrors the background-job model used by `picture_agent` (and the upstream
//! poll + workspace-transfer + billing flow of `digital_human`): a submit
//! endpoint inserts a `background_tasks` row, returns the local task id
//! immediately, and a `tokio::spawn`ed worker drives submit → poll → store →
//! bill while updating progress. Video output URLs are valid for only 1 hour,
//! so finished outputs are transferred into the workspace immediately and only
//! the persisted local asset is exposed.
use crate::{
auth::jwt::validate_jwt,
db::models::background_task::BackgroundTask,
db::models::downstream_api_key::DownstreamApiKey,
error::{AppError, AppResult},
services::billing::{ComfyTaskBillingEvent, UsageTracker},
services::jimeng::poller::poll_until_done,
services::jimeng::types::{ImageRequest, JimengCapability, JimengOutput, VideoRequest},
services::jimeng::{build_image_body, build_video_body, JimengError, JimengProvider},
services::workspace::{record_base64_asset, record_remote_asset, WorkspaceAssetInput},
AppState,
};
use axum::{
extract::{Path, Query, State},
http::{header, HeaderMap},
routing::{get, post},
Json, Router,
};
use serde::Deserialize;
use serde_json::{json, Value};
use std::time::Duration;
const TASK_TYPE_IMAGE: &str = "jimeng-image-generation";
const TASK_TYPE_VIDEO: &str = "jimeng-video-generation";
const PROGRESS_TOTAL: i64 = 4;
#[derive(Debug, Clone)]
struct JimengAuth {
user_id: Option<i64>,
role: String,
key_id: Option<i64>,
key_name: String,
}
pub fn routes() -> Router<AppState> {
Router::new()
.route("/api/jimeng/status", get(status_handler))
.route("/api/jimeng/image/jobs", post(create_image_job))
.route("/api/jimeng/video/jobs", post(create_video_job))
.route("/api/jimeng/jobs", get(list_jobs))
.route("/api/jimeng/jobs/{id}", get(get_job))
.route("/api/jimeng/jobs/{id}/cancel", post(cancel_job))
}
fn extract_bearer_token(headers: &HeaderMap) -> Option<&str> {
let auth_header = headers
.get(header::AUTHORIZATION)
.and_then(|value| value.to_str().ok())
.unwrap_or_default();
auth_header
.strip_prefix("Bearer ")
.or_else(|| auth_header.strip_prefix("bearer "))
.map(str::trim)
.filter(|token| !token.is_empty())
}
async fn user_downstream_key(state: &AppState, user_id: i64) -> Option<DownstreamApiKey> {
sqlx::query_as::<_, DownstreamApiKey>(
"SELECT *
FROM downstream_api_keys
WHERE user_id = ?
AND COALESCE(enabled, 1) = 1
AND (
expires_at IS NULL
OR expires_at = ''
OR datetime(expires_at) > datetime('now')
)
ORDER BY updated_at DESC, created_at DESC, id DESC
LIMIT 1",
)
.bind(user_id)
.fetch_optional(&*state.db)
.await
.ok()
.flatten()
}
/// Resolve the caller the same way `picture_agent` does: the global proxy /
/// admin token maps to an admin caller, otherwise a JWT identifies a user (and
/// their downstream key id is recorded for billing).
async fn resolve_auth(headers: &HeaderMap, state: &AppState) -> Result<JimengAuth, AppError> {
let token = extract_bearer_token(headers)
.ok_or_else(|| AppError::Unauthorized("Missing Authorization header".into()))?;
let proxy_token = state.runtime_proxy_token();
if token == state.config.auth.auth_token || token == proxy_token {
return Ok(JimengAuth {
user_id: None,
role: "admin".to_string(),
key_id: None,
key_name: "global-proxy".to_string(),
});
}
let claims = validate_jwt(&state.config.auth.jwt_secret, token)
.map_err(|_| AppError::Unauthorized("Invalid or expired token".into()))?;
if claims.role != "admin" && claims.role != "user" {
return Err(AppError::Forbidden("Forbidden".into()));
}
let user_id = claims
.sub
.parse::<i64>()
.map_err(|_| AppError::Unauthorized("Invalid user ID in token".into()))?;
match user_downstream_key(state, user_id).await {
Some(key) => Ok(JimengAuth {
user_id: Some(user_id),
role: claims.role,
key_id: Some(key.id),
key_name: key.name,
}),
None if claims.role == "admin" => Ok(JimengAuth {
user_id: Some(user_id),
role: claims.role,
key_id: None,
key_name: "global-proxy".to_string(),
}),
None => Err(AppError::Forbidden(
"当前用户没有可用的 API Key,请先在用户中心创建或启用 API Key".into(),
)),
}
}
fn provider_or_unavailable(state: &AppState) -> AppResult<JimengProvider> {
JimengProvider::from_config(&state.config.jimeng).map_err(|err| match err {
JimengError::Disabled => {
AppError::ServiceUnavailable("即梦能力未启用,请在配置中开启 [jimeng]".into())
}
other => AppError::Internal(anyhow::anyhow!(other.to_string())),
})
}
// PLACEHOLDER_REST
async fn status_handler(State(state): State<AppState>, headers: HeaderMap) -> AppResult<Json<Value>> {
let auth = resolve_auth(&headers, &state).await?;
let cfg = &state.config.jimeng;
Ok(Json(json!({
"success": true,
"enabled": cfg.enabled,
"role": auth.role,
"userId": auth.user_id,
"keyName": auth.key_name,
"capabilities": {
"image": { "reqKey": JimengCapability::Image.req_key() },
"video": { "reqKey": JimengCapability::Video.req_key() },
},
})))
}
#[derive(Debug, Deserialize)]
struct JobListQuery {
limit: Option<i64>,
}
/// A task is visible to the caller when they are the global admin, or when the
/// task's stored `userId` matches the caller's id.
fn job_visible_to_auth(task: &BackgroundTask, auth: &JimengAuth) -> bool {
if auth.role == "admin" && auth.user_id.is_none() {
return true;
}
let Some(raw) = task.result_json.as_deref() else {
return false;
};
let Ok(value) = serde_json::from_str::<Value>(raw) else {
return false;
};
value.get("userId").and_then(Value::as_i64) == auth.user_id
}
fn task_to_json(task: BackgroundTask) -> Value {
let result = task
.result_json
.as_deref()
.and_then(|raw| serde_json::from_str::<Value>(raw).ok())
.unwrap_or_else(|| json!({}));
json!({
"id": task.id,
"type": task.task_type,
"title": task.title,
"status": task.status,
"message": task.message,
"error": task.error,
"progress": task.progress.unwrap_or(0),
"progressTotal": task.progress_total.unwrap_or(PROGRESS_TOTAL),
"result": result,
"createdAt": task.created_at,
"updatedAt": task.updated_at,
"startedAt": task.started_at,
"finishedAt": task.finished_at,
})
}
async fn load_job(state: &AppState, id: &str, auth: &JimengAuth) -> AppResult<BackgroundTask> {
let task: BackgroundTask = sqlx::query_as(
"SELECT * FROM background_tasks
WHERE id = ? AND task_type IN (?, ?) LIMIT 1",
)
.bind(id)
.bind(TASK_TYPE_IMAGE)
.bind(TASK_TYPE_VIDEO)
.fetch_optional(&*state.db)
.await?
.ok_or_else(|| AppError::NotFound("Jimeng job not found".into()))?;
if !job_visible_to_auth(&task, auth) {
return Err(AppError::NotFound("Jimeng job not found".into()));
}
Ok(task)
}
async fn update_job_progress(state: &AppState, id: &str, progress: i64, message: &str) {
let _ = sqlx::query(
"UPDATE background_tasks
SET progress = ?, message = ?, updated_at = datetime('now')
WHERE id = ? AND status IN ('pending', 'running')",
)
.bind(progress)
.bind(message)
.bind(id)
.execute(&*state.db)
.await;
}
async fn is_job_cancelled(state: &AppState, id: &str) -> bool {
let status: Option<String> =
sqlx::query_scalar("SELECT status FROM background_tasks WHERE id = ? LIMIT 1")
.bind(id)
.fetch_optional(&*state.db)
.await
.ok()
.flatten();
status.as_deref() == Some("cancelled")
}
// PLACEHOLDER_HANDLERS
#[derive(Debug, Deserialize)]
struct ImageJobRequest {
prompt: String,
#[serde(default)]
width: Option<i64>,
#[serde(default)]
height: Option<i64>,
#[serde(default)]
scale: Option<f64>,
#[serde(default)]
force_single: Option<bool>,
}
#[derive(Debug, Deserialize)]
struct VideoJobRequest {
#[serde(default)]
prompt: Option<String>,
#[serde(default)]
image_urls: Vec<String>,
#[serde(default = "default_video_frames")]
frames: i64,
#[serde(default)]
aspect_ratio: Option<String>,
#[serde(default)]
seed: Option<i64>,
}
fn default_video_frames() -> i64 {
121
}
async fn create_image_job(
State(state): State<AppState>,
headers: HeaderMap,
Json(req): Json<ImageJobRequest>,
) -> AppResult<Json<Value>> {
let auth = resolve_auth(&headers, &state).await?;
let provider = provider_or_unavailable(&state)?;
let image_req = ImageRequest {
prompt: req.prompt,
width: req.width,
height: req.height,
scale: req.scale,
force_single: req.force_single,
};
// Validate eagerly so the caller gets a 400 instead of a failed job.
let body = build_image_body(&image_req).map_err(invalid_to_apperr)?;
let title = image_req.prompt.chars().take(30).collect::<String>();
let task_id = spawn_job(
&state,
&auth,
provider,
JimengCapability::Image,
body,
title,
json!({ "prompt": image_req.prompt }),
)
.await?;
let task = load_job(&state, &task_id, &auth).await?;
Ok(Json(json!({ "success": true, "job": task_to_json(task) })))
}
async fn create_video_job(
State(state): State<AppState>,
headers: HeaderMap,
Json(req): Json<VideoJobRequest>,
) -> AppResult<Json<Value>> {
let auth = resolve_auth(&headers, &state).await?;
let provider = provider_or_unavailable(&state)?;
let video_req = VideoRequest {
prompt: req.prompt,
image_urls: req.image_urls,
frames: req.frames,
aspect_ratio: req.aspect_ratio,
seed: req.seed,
};
let body = build_video_body(&video_req).map_err(invalid_to_apperr)?;
let title = video_req
.prompt
.as_deref()
.unwrap_or("图生视频")
.chars()
.take(30)
.collect::<String>();
let task_id = spawn_job(
&state,
&auth,
provider,
JimengCapability::Video,
body,
title,
json!({
"prompt": video_req.prompt,
"imageUrls": video_req.image_urls,
"frames": video_req.frames,
}),
)
.await?;
let task = load_job(&state, &task_id, &auth).await?;
Ok(Json(json!({ "success": true, "job": task_to_json(task) })))
}
fn invalid_to_apperr(err: JimengError) -> AppError {
match err {
JimengError::Invalid(message) => AppError::BadRequest(message),
other => AppError::BadRequest(other.to_string()),
}
}
async fn list_jobs(
State(state): State<AppState>,
headers: HeaderMap,
Query(query): Query<JobListQuery>,
) -> AppResult<Json<Value>> {
let auth = resolve_auth(&headers, &state).await?;
let limit = query.limit.unwrap_or(20).clamp(1, 100);
let tasks: Vec<BackgroundTask> = sqlx::query_as(
"SELECT * FROM background_tasks
WHERE task_type IN (?, ?)
ORDER BY created_at DESC
LIMIT ?",
)
.bind(TASK_TYPE_IMAGE)
.bind(TASK_TYPE_VIDEO)
.bind(limit * 3)
.fetch_all(&*state.db)
.await?;
let jobs: Vec<Value> = tasks
.into_iter()
.filter(|task| job_visible_to_auth(task, &auth))
.take(limit as usize)
.map(task_to_json)
.collect();
Ok(Json(json!({ "success": true, "jobs": jobs })))
}
async fn get_job(
State(state): State<AppState>,
headers: HeaderMap,
Path(id): Path<String>,
) -> AppResult<Json<Value>> {
let auth = resolve_auth(&headers, &state).await?;
let task = load_job(&state, &id, &auth).await?;
Ok(Json(json!({ "success": true, "job": task_to_json(task) })))
}
async fn cancel_job(
State(state): State<AppState>,
headers: HeaderMap,
Path(id): Path<String>,
) -> AppResult<Json<Value>> {
let auth = resolve_auth(&headers, &state).await?;
let task = load_job(&state, &id, &auth).await?;
if task.status == "pending" || task.status == "running" {
sqlx::query(
"UPDATE background_tasks
SET status = 'cancelled', message = '已取消', error = NULL,
updated_at = datetime('now'), finished_at = datetime('now')
WHERE id = ? AND status IN ('pending', 'running')",
)
.bind(&id)
.execute(&*state.db)
.await?;
}
let task = load_job(&state, &id, &auth).await?;
Ok(Json(json!({ "success": true, "job": task_to_json(task) })))
}
// PLACEHOLDER_WORKER
/// Insert the `background_tasks` row, spawn the worker, and return the local id.
async fn spawn_job(
state: &AppState,
auth: &JimengAuth,
provider: JimengProvider,
capability: JimengCapability,
body: Value,
title: String,
request_snapshot: Value,
) -> AppResult<String> {
let task_id = format!("{}_{}", capability.id_prefix(), uuid::Uuid::new_v4().simple());
let task_type = capability.task_type();
let initial_result = json!({
"userId": auth.user_id,
"role": auth.role,
"keyId": auth.key_id,
"keyName": auth.key_name,
"capability": capability.req_key(),
"request": request_snapshot,
});
sqlx::query(
"INSERT INTO background_tasks
(id, task_type, title, status, message, dedupe_key, result_json, progress, progress_total, started_at)
VALUES (?, ?, ?, 'running', '任务已提交,等待火山', ?, ?, 1, ?, datetime('now'))",
)
.bind(&task_id)
.bind(task_type)
.bind(if title.trim().is_empty() { "即梦生成任务" } else { title.trim() })
.bind(format!(
"jimeng:{}:{}",
capability.req_key(),
auth.user_id.map(|id| id.to_string()).unwrap_or_else(|| "admin".into())
))
.bind(initial_result.to_string())
.bind(PROGRESS_TOTAL)
.execute(&*state.db)
.await?;
let state_for_task = state.clone();
let auth_for_task = auth.clone();
let task_id_spawn = task_id.clone();
tokio::spawn(async move {
run_worker(
state_for_task,
auth_for_task,
provider,
capability,
body,
task_id_spawn,
initial_result,
)
.await;
});
Ok(task_id)
}
/// Submit → poll → transfer-to-workspace → bill. Terminal status and billing
/// are written here; progress advances 1→4 as the task moves through stages.
async fn run_worker(
state: AppState,
auth: JimengAuth,
provider: JimengProvider,
capability: JimengCapability,
body: Value,
task_id: String,
base_result: Value,
) {
let model = capability.req_key().to_string();
let outcome = drive(&state, &provider, capability, &body, &task_id).await;
match outcome {
Ok(output) => {
update_job_progress(&state, &task_id, 3, "生成完成,正在转存到工作空间").await;
let assets = store_outputs(&state, &auth, capability, &output).await;
let mut final_result = base_result.clone();
if let Some(obj) = final_result.as_object_mut() {
obj.insert("assets".into(), json!(assets));
}
let _ = sqlx::query(
"UPDATE background_tasks
SET status = 'succeeded', message = '生成完成', result_json = ?, progress = 4,
updated_at = datetime('now'), finished_at = datetime('now')
WHERE id = ? AND status IN ('pending', 'running')",
)
.bind(final_result.to_string())
.bind(&task_id)
.execute(&*state.db)
.await;
record_billing(&state, &auth, capability, &model, true, None).await;
}
Err(JimengError::Cancelled) => {
// Cancellation already set the terminal status; nothing to bill.
}
Err(err) => {
let _ = sqlx::query(
"UPDATE background_tasks
SET status = 'failed', message = '生成失败', error = ?,
updated_at = datetime('now'), finished_at = datetime('now')
WHERE id = ? AND status IN ('pending', 'running')",
)
.bind(err.to_string())
.bind(&task_id)
.execute(&*state.db)
.await;
record_billing(&state, &auth, capability, &model, false, Some(err.to_string())).await;
}
}
}
/// Submit to Volcengine then poll to completion, honoring cancellation.
async fn drive(
state: &AppState,
provider: &JimengProvider,
capability: JimengCapability,
body: &Value,
task_id: &str,
) -> Result<JimengOutput, JimengError> {
let volc_task_id = provider.submit(&state.http, capability, body).await?;
update_job_progress(state, task_id, 2, "已提交火山,正在生成").await;
// Record the upstream task id so it survives a poll restart / inspection.
let _ = sqlx::query(
"UPDATE background_tasks
SET result_json = json_set(COALESCE(result_json, '{}'), '$.volcTaskId', ?),
updated_at = datetime('now')
WHERE id = ?",
)
.bind(&volc_task_id)
.bind(task_id)
.execute(&*state.db)
.await;
let cfg = &state.config.jimeng;
let (interval, timeout) = match capability {
JimengCapability::Image => (
Duration::from_secs(cfg.image_poll_interval_secs.max(1)),
Duration::from_secs(cfg.image_timeout_secs.max(1)),
),
JimengCapability::Video => (
Duration::from_secs(cfg.video_poll_interval_secs.max(1)),
Duration::from_secs(cfg.video_timeout_secs.max(1)),
),
};
poll_until_done(
capability,
interval,
timeout,
|| provider.get_result(&state.http, capability, &volc_task_id),
|| is_job_cancelled(state, task_id),
|d| tokio::time::sleep(d),
)
.await
}
/// Transfer finished outputs into the workspace and return their descriptors.
/// Video URLs expire after 1h, so this runs immediately on completion and we
/// only ever expose the persisted local asset.
async fn store_outputs(
state: &AppState,
auth: &JimengAuth,
capability: JimengCapability,
output: &JimengOutput,
) -> Vec<Value> {
let model = capability.req_key().to_string();
let mut assets = Vec::new();
match output {
JimengOutput::Image { urls, base64 } => {
for url in urls {
let input = asset_input(auth, capability, &model, Some(url.clone()));
let stored = record_remote_asset(state, input).await;
assets.push(asset_to_json(stored));
}
for data in base64 {
let input = asset_input(auth, capability, &model, None);
match record_base64_asset(state, input, data).await {
Ok(stored) => assets.push(asset_to_json(Ok(stored))),
Err(err) => assets.push(json!({ "status": "failed", "error": err.to_string() })),
}
}
}
JimengOutput::Video { url } => {
let input = asset_input(auth, capability, &model, Some(url.clone()));
let stored = record_remote_asset(state, input).await;
assets.push(asset_to_json(stored));
}
}
assets
}
fn asset_input(
auth: &JimengAuth,
capability: JimengCapability,
model: &str,
remote_url: Option<String>,
) -> WorkspaceAssetInput {
WorkspaceAssetInput {
user_id: auth.user_id,
source: "jimeng".to_string(),
kind: capability.asset_kind().to_string(),
title: None,
prompt: None,
model: Some(model.to_string()),
remote_url,
metadata: json!({ "jimeng": { "reqKey": model } }),
}
}
fn asset_to_json(
stored: Result<crate::services::workspace::StoredWorkspaceAsset, AppError>,
) -> Value {
match stored {
Ok(asset) => json!({
"publicId": asset.public_id,
"status": asset.status,
"fileUrl": format!("/api/user/workspace/assets/{}/file", asset.public_id),
}),
Err(err) => json!({ "status": "failed", "error": err.to_string() }),
}
}
async fn record_billing(
state: &AppState,
auth: &JimengAuth,
capability: JimengCapability,
model: &str,
success: bool,
error_message: Option<String>,
) {
let cfg = &state.config.jimeng;
let cost = match capability {
JimengCapability::Image => cfg.image_cost,
JimengCapability::Video => cfg.video_cost,
};
let event = ComfyTaskBillingEvent {
key_id: auth.key_id,
model: model.to_string(),
task_type: capability.task_type(),
client_family: "jimeng",
client_app_id: "jimeng",
client_app_name: "Jimeng",
prompt_tokens: 0,
completion_tokens: 0,
cost,
content_summary: None,
request_content: None,
response_content: None,
error_message,
};
let tracker = UsageTracker::new(state.db.clone());
let result = if success {
tracker.record_comfy_task_success(event).await
} else {
tracker.record_comfy_task_failure(event).await
};
if let Err(err) = result {
tracing::warn!(%err, "jimeng billing record failed");
}
}
+2
View File
@@ -31,6 +31,7 @@ mod comfyui_agent;
mod digital_human; mod digital_human;
mod events; mod events;
mod health; mod health;
mod jimeng;
mod money_printer; mod money_printer;
mod oauth; mod oauth;
mod organizations; mod organizations;
@@ -61,6 +62,7 @@ pub fn build_router(state: AppState) -> Router {
.merge(agent::routes()) .merge(agent::routes())
.merge(digital_human::routes()) .merge(digital_human::routes())
.merge(comfyui::routes()) .merge(comfyui::routes())
.merge(jimeng::routes())
.merge(crate::workflow::routes::routes()) .merge(crate::workflow::routes::routes())
.route("/comfy", any(removed_comfy_workbench)) .route("/comfy", any(removed_comfy_workbench))
.route("/comfy/", any(removed_comfy_workbench)) .route("/comfy/", any(removed_comfy_workbench))
+1 -1
View File
@@ -121,7 +121,7 @@ pub fn classify(code: i64, cap: JimengCapability) -> CodeClass {
// Post-generation text moderation failure: not retryable. // Post-generation text moderation failure: not retryable.
50512 => CodeClass::Fatal, 50512 => CodeClass::Fatal,
// Moderation / copyright service exceptions (upstream): not retryable. // Moderation / copyright service exceptions (upstream): not retryable.
50520 | 50521 | 50522 => CodeClass::Fatal, 50520..=50522 => CodeClass::Fatal,
// Output post-moderation failures (image / copyright): retryable. // Output post-moderation failures (image / copyright): retryable.
50511 | 50519 => CodeClass::Retryable, 50511 | 50519 => CodeClass::Retryable,
// Output video / audio post-moderation failure (video only): retryable. // Output video / audio post-moderation failure (video only): retryable.