Genarrative/server-rs/crates/api-server/src/editor_generation_queue.rs

use axum::http::StatusCode;
use serde::Serialize;
use serde_json::{Value, json};
use shared_contracts::external_generation::{
    ExternalGenerationJobStatus, ExternalGenerationJobStatusRecord,
};
use shared_kernel::{build_prefixed_uuid_id, offset_datetime_to_unix_micros};
use spacetime_client::{ExternalGenerationJobEnqueueRecordInput, ExternalGenerationJobRecord};

use crate::{http_error::AppError, request_context::RequestContext, state::AppState};

pub(crate) const EDITOR_IMAGE_GENERATION_JOB_KIND: &str = "editor_image_generation";
pub(crate) const EDITOR_IMAGE_EDIT_JOB_KIND: &str = "editor_image_edit";
pub(crate) const EDITOR_BACKGROUND_REMOVAL_JOB_KIND: &str = "editor_background_removal";
pub(crate) const EDITOR_ICON_SPRITESHEET_GENERATION_JOB_KIND: &str =
    "editor_icon_spritesheet_generation";
pub(crate) const EDITOR_UI_DESIGN_ASSET_EXTRACTION_JOB_KIND: &str =
    "editor_ui_design_asset_extraction";
pub(crate) const EDITOR_CHARACTER_ANIMATION_GENERATION_JOB_KIND: &str =
    "editor_character_animation_generation";
pub(crate) const EDITOR_VIDEO_GENERATION_JOB_KIND: &str = "editor_video_generation";
pub(crate) const EDITOR_SOUND_EFFECT_GENERATION_JOB_KIND: &str = "editor_sound_effect_generation";
pub(crate) const EDITOR_BACKGROUND_MUSIC_GENERATION_JOB_KIND: &str =
    "editor_background_music_generation";

pub(crate) const EDITOR_GENERATION_QUEUE_SOURCE_MODULE: &str = "editor-canvas";
const EDITOR_GENERATION_QUEUE_PROVIDER: &str = "editor-generation-worker";
const MAX_EDITOR_GENERATION_JOB_PAYLOAD_BYTES: usize = 512 * 1024;

#[derive(Debug, Serialize)]
#[serde(rename_all = "camelCase")]
pub(crate) struct EditorGenerationQueuedResponse {
    pub(crate) queue_state: ExternalGenerationJobStatusRecord,
}

pub(crate) async fn enqueue_editor_generation_job<T>(
    state: &AppState,
    _request_context: &RequestContext,
    owner_user_id: &str,
    job_kind: &str,
    source_entity_id: impl Into<String>,
    request_label: impl Into<String>,
    price_mud_points: u64,
    payload: &T,
) -> Result<ExternalGenerationJobRecord, AppError>
where
    T: Serialize,
{
    let job_id = build_prefixed_uuid_id("task-");
    enqueue_editor_generation_job_with_identity(
        state,
        owner_user_id,
        job_kind,
        source_entity_id,
        request_label,
        price_mud_points,
        payload,
        job_id.clone(),
        format!("editor-canvas:{job_kind}:{job_id}"),
    )
    .await
}

#[allow(clippy::too_many_arguments)]
pub(crate) async fn enqueue_editor_generation_job_with_identity<T>(
    state: &AppState,
    owner_user_id: &str,
    job_kind: &str,
    source_entity_id: impl Into<String>,
    request_label: impl Into<String>,
    price_mud_points: u64,
    payload: &T,
    job_id: String,
    dedupe_key: String,
) -> Result<ExternalGenerationJobRecord, AppError>
where
    T: Serialize,
{
    let request_payload_json = serialize_editor_generation_job_payload(payload)?;
    let now_micros = current_utc_micros();
    state
        .spacetime_client()
        .enqueue_external_generation_job(ExternalGenerationJobEnqueueRecordInput {
            dedupe_key,
            job_id,
            job_kind: job_kind.to_string(),
            owner_user_id: owner_user_id.to_string(),
            source_module: EDITOR_GENERATION_QUEUE_SOURCE_MODULE.to_string(),
            source_entity_id: source_entity_id.into(),
            request_label: request_label.into(),
            request_payload_json,
            max_attempts: 1,
            available_at_micros: now_micros,
            created_at_micros: now_micros,
            price_mud_points,
        })
        .await
        .map_err(|error| {
            AppError::from_status(StatusCode::BAD_GATEWAY).with_details(json!({
                "provider": EDITOR_GENERATION_QUEUE_PROVIDER,
                "message": error.to_string(),
            }))
        })
}

fn serialize_editor_generation_job_payload<T>(payload: &T) -> Result<String, AppError>
where
    T: Serialize + ?Sized,
{
    let payload_value = serde_json::to_value(payload).map_err(payload_serialization_error)?;
    if contains_inline_media_reference(&payload_value) {
        return Err(
            AppError::from_status(StatusCode::BAD_REQUEST).with_details(json!({
                "provider": EDITOR_GENERATION_QUEUE_PROVIDER,
                "message": "编辑器生成任务参数禁止包含 data: 或 blob: 内联媒体引用，请先将媒体上传到对象存储并改传 objectKey 或 resourceId。",
            })),
        );
    }

    let request_payload_json =
        serde_json::to_string(&payload_value).map_err(payload_serialization_error)?;
    let payload_bytes = request_payload_json.len();
    if payload_bytes > MAX_EDITOR_GENERATION_JOB_PAYLOAD_BYTES {
        return Err(
            AppError::from_status(StatusCode::PAYLOAD_TOO_LARGE).with_details(json!({
                "provider": EDITOR_GENERATION_QUEUE_PROVIDER,
                "message": format!(
                    "编辑器生成任务 JSON 大小为 {payload_bytes} 字节，超过持久化上限 {MAX_EDITOR_GENERATION_JOB_PAYLOAD_BYTES} 字节；请移除冗余数据并改传 objectKey 或 resourceId。"
                ),
                "actualBytes": payload_bytes,
                "maxBytes": MAX_EDITOR_GENERATION_JOB_PAYLOAD_BYTES,
            })),
        );
    }

    Ok(request_payload_json)
}

fn payload_serialization_error(error: serde_json::Error) -> AppError {
    AppError::from_status(StatusCode::INTERNAL_SERVER_ERROR).with_details(json!({
        "provider": EDITOR_GENERATION_QUEUE_PROVIDER,
        "message": format!("编辑器 worker 任务参数序列化失败：{error}"),
    }))
}

fn contains_inline_media_reference(value: &Value) -> bool {
    match value {
        Value::String(value) => is_inline_media_reference(value),
        Value::Array(values) => values.iter().any(contains_inline_media_reference),
        Value::Object(values) => values.iter().any(|(key, value)| {
            is_inline_media_reference(key) || contains_inline_media_reference(value)
        }),
        Value::Null | Value::Bool(_) | Value::Number(_) => false,
    }
}

fn is_inline_media_reference(value: &str) -> bool {
    let prefix = value.trim_start().as_bytes().get(..5);
    prefix.is_some_and(|prefix| {
        prefix.eq_ignore_ascii_case(b"data:") || prefix.eq_ignore_ascii_case(b"blob:")
    })
}

pub(crate) fn editor_generation_queue_state(
    job: ExternalGenerationJobRecord,
) -> ExternalGenerationJobStatusRecord {
    ExternalGenerationJobStatusRecord {
        operation_id: job.job_id,
        status: ExternalGenerationJobStatus::Queued,
        phase_label: job.request_label,
        phase_detail: "排队中。".to_string(),
        progress: 8,
        error: job.last_error_message,
        updated_at_micros: job.updated_at_micros,
    }
}

pub(crate) fn editor_generation_source_entity_id(
    project_id: Option<&str>,
    fallback: &str,
) -> String {
    project_id
        .map(str::trim)
        .filter(|value| !value.is_empty())
        .unwrap_or(fallback)
        .to_string()
}

fn current_utc_micros() -> i64 {
    offset_datetime_to_unix_micros(time::OffsetDateTime::now_utc())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn serialize_payload_accepts_persistable_media_references() {
        let payload = json!({
            "sourceImageObjectKey": "users/user-1/editor/source.png",
            "resourceId": "resource-1",
            "nested": [{ "prompt": "保留 data 与 blob 这两个普通单词" }],
        });

        let serialized = serialize_editor_generation_job_payload(&payload)
            .expect("objectKey 和 resourceId 应允许进入持久任务 JSON");

        assert_eq!(
            serde_json::from_str::<serde_json::Value>(&serialized).expect("应生成有效 JSON"),
            payload
        );
    }

    #[test]
    fn serialize_payload_rejects_nested_data_url_case_insensitively() {
        let payload = json!({
            "input": {
                "references": [
                    { "url": "  \nDaTa:image/png;base64,AAAA" }
                ]
            }
        });

        let error = serialize_editor_generation_job_payload(&payload)
            .expect_err("任意层级的 Data URL 都必须被拒绝");

        assert_eq!(error.status_code(), StatusCode::BAD_REQUEST);
        assert!(error.body_text().contains("禁止包含 data: 或 blob:"));
    }

    #[test]
    fn serialize_payload_rejects_nested_blob_url_case_insensitively() {
        let payload = json!({
            "input": [{ "source": { "url": "\tBLOB:https://example.test/id" } }]
        });

        let error = serialize_editor_generation_job_payload(&payload)
            .expect_err("任意层级的 Blob URL 都必须被拒绝");

        assert_eq!(error.status_code(), StatusCode::BAD_REQUEST);
        assert!(error.body_text().contains("禁止包含 data: 或 blob:"));
    }

    #[test]
    fn serialize_payload_rejects_json_larger_than_persistence_limit() {
        let payload = json!({
            "prompt": "x".repeat(MAX_EDITOR_GENERATION_JOB_PAYLOAD_BYTES),
        });

        let error = serialize_editor_generation_job_payload(&payload)
            .expect_err("超过上限的任务 JSON 必须被拒绝");

        assert_eq!(error.status_code().as_u16(), 413);
        assert!(error.body_text().contains("超过持久化上限"));
    }
}