feat(jimeng): add image/video generation routes with background jobs, workspace transfer and billing
This commit is contained in:
@@ -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");
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
@@ -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))
|
||||||
|
|||||||
@@ -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.
|
||||||
|
|||||||
Reference in New Issue
Block a user