From 764a6e1c8d5cb9a67f9055630d73342486b4d90e Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Tue, 2 Jun 2026 05:47:24 -0400 Subject: [PATCH] add correct job type --- .../image/insert_db_job/insert_fal_job.rs | 59 +++++++++ .../insert_db_job/insert_seedance2pro_jobs.rs | 115 ++++++++++++++++++ .../generate/image/insert_db_job/mod.rs | 3 + .../image/insert_db_job/shared_job_args.rs | 20 +++ .../endpoints/omni_gen/generate/image/mod.rs | 1 + .../image/omni_gen_image_generate_handler.rs | 106 +++++++++------- .../generate/image/pipeline_result.rs | 7 ++ .../image/pipeline_v2/run_pipeline_v2.rs | 10 +- 8 files changed, 272 insertions(+), 49 deletions(-) create mode 100644 crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/insert_db_job/insert_fal_job.rs create mode 100644 crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/insert_db_job/insert_seedance2pro_jobs.rs create mode 100644 crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/insert_db_job/mod.rs create mode 100644 crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/insert_db_job/shared_job_args.rs diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/insert_db_job/insert_fal_job.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/insert_db_job/insert_fal_job.rs new file mode 100644 index 0000000000..19df43757e --- /dev/null +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/insert_db_job/insert_fal_job.rs @@ -0,0 +1,59 @@ +use log::warn; + +use enums::common::visibility::Visibility; +use mysql_queries::queries::generic_inference::api_providers::fal::insert_generic_inference_job_for_fal_queue::FalCategory; +use mysql_queries::queries::generic_inference::api_providers::fal::insert_generic_inference_job_for_fal_queue_with_apriori_job_token::{ + insert_generic_inference_job_for_fal_queue_with_apriori_job_token, + InsertGenericInferenceForFalWithAprioriJobTokenArgs, +}; +use tokens::tokens::generic_inference_jobs::InferenceJobToken; + +use crate::http_server::common_responses::common_web_error::CommonWebError; + +use super::shared_job_args::SharedJobArgs; + +pub struct InsertFalJobArgs<'a, 'tx> { + pub external_job_id: &'a str, + pub shared: SharedJobArgs<'a, 'tx>, +} + +/// Insert a single Fal-queue inference job row for an image generation. +/// Mirrors `omni_gen::generate::video::insert_db_job::insert_fal_job::insert_fal_job` +/// but tags the row with `FalCategory::ImageGeneration` so the worker +/// routes it to the image-generation queue. +pub async fn insert_fal_job(args: InsertFalJobArgs<'_, '_>) -> Result { + let InsertFalJobArgs { + external_job_id, + shared, + } = args; + + let db_result = insert_generic_inference_job_for_fal_queue_with_apriori_job_token( + InsertGenericInferenceForFalWithAprioriJobTokenArgs { + apriori_job_token: shared.apriori_job_token, + uuid_idempotency_token: shared.idempotency_token, + maybe_external_third_party_id: external_job_id, + fal_category: FalCategory::ImageGeneration, + maybe_model_type: shared.maybe_model_type, + maybe_prompt_token: shared.maybe_prompt_token, + maybe_creator_user_token: Some(shared.user_token), + maybe_avt_token: shared.maybe_avt_token, + creator_ip_address: shared.ip_address, + creator_set_visibility: Visibility::Public, + maybe_debug_log_event_token: shared.maybe_debug_log_event_token, + mysql_executor: &mut **shared.transaction, + maybe_inference_args: None, + starting_job_status_override: None, + maybe_frontend_failure_category: None, + maybe_failure_reason: None, + phantom: Default::default(), + } + ).await; + + match db_result { + Ok(token) => Ok(token), + Err(err) => { + warn!("Error inserting fal inference job (image): {:?}", err); + Err(CommonWebError::from_error(err)) + } + } +} diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/insert_db_job/insert_seedance2pro_jobs.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/insert_db_job/insert_seedance2pro_jobs.rs new file mode 100644 index 0000000000..00c6ca5c9b --- /dev/null +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/insert_db_job/insert_seedance2pro_jobs.rs @@ -0,0 +1,115 @@ +use log::{error, warn}; + +use enums::common::visibility::Visibility; +use mysql_queries::queries::generic_inference::api_providers::seedance2pro::insert_generic_inference_job_for_seedance2pro_queue_with_apriori_job_token::{ + insert_generic_inference_job_for_seedance2pro_queue_with_apriori_job_token, + InsertGenericInferenceForSeedance2ProWithAprioriJobTokenArgs, KinoviVersion, +}; +use tokens::tokens::generic_inference_jobs::InferenceJobToken; +use tokens::tokens::wallet_ledger_entries::WalletLedgerEntryToken; +use uuid_utils::uuid::generate_random_uuid; + +use crate::http_server::common_responses::common_web_error::CommonWebError; + +use super::shared_job_args::SharedJobArgs; + +pub struct InsertSeedance2proJobsArgs<'a, 'tx> { + /// The order_id for the primary job (always inserted first). + pub primary_order_id: &'a str, + + /// Additional order_ids returned by Kinovi when batch_count > 1. Each + /// becomes its own inference job row. + pub maybe_additional_order_ids: Option<&'a [String]>, + + /// Wallet ledger entry from the up-front deduction. Stamped on each + /// inserted row so the worker can issue a refund on failure. + pub maybe_wallet_ledger_entry_token: Option<&'a WalletLedgerEntryToken>, + + /// Which Kinovi account fulfilled the request — Volcengine, BytePlus, + /// or BytePlusUltra. For Midjourney image generation today this is + /// always Volcengine, but the helper stays version-agnostic so a future + /// account-picking layer can use it unchanged. + pub kinovi_version: KinoviVersion, + + pub shared: SharedJobArgs<'a, 'tx>, +} + +pub struct InsertSeedance2proJobsResult { + pub primary_job_token: InferenceJobToken, + pub all_job_tokens: Vec, +} + +/// Insert one Seedance2Pro/Kinovi inference job per returned order_id. +/// Mirrors the video-side helper at +/// `omni_gen::generate::video::insert_db_job::insert_seedance2pro_jobs`. +pub async fn insert_seedance2pro_jobs(args: InsertSeedance2proJobsArgs<'_, '_>) -> Result { + let InsertSeedance2proJobsArgs { + primary_order_id, + maybe_additional_order_ids, + maybe_wallet_ledger_entry_token, + kinovi_version, + mut shared, + } = args; + + // Build a deduplicated list with primary_order_id first. + let mut all_order_ids = vec![primary_order_id.to_string()]; + + if let Some(additional) = maybe_additional_order_ids { + for id in additional { + if id != primary_order_id { + all_order_ids.push(id.clone()); + } + } + } + + let mut all_job_tokens: Vec = Vec::with_capacity(all_order_ids.len()); + + for (i, order_id) in all_order_ids.iter().enumerate() { + let job_token = if i == 0 { shared.apriori_job_token.clone() } else { InferenceJobToken::generate() }; + + // The idempotency_token column has a max length too tight for the + // `{idempotency}-batch-{i}` suffix pattern used elsewhere; mint a + // fresh UUID for each additional row instead. The primary row still + // wants a deterministic value, but historic behavior on the video + // side has been to mint UUIDs for all rows, which we mirror here. + let idempotency_str = generate_random_uuid(); + + let db_result = insert_generic_inference_job_for_seedance2pro_queue_with_apriori_job_token( + InsertGenericInferenceForSeedance2ProWithAprioriJobTokenArgs { + kinovi_version, + apriori_job_token: &job_token, + uuid_idempotency_token: &idempotency_str, + maybe_external_third_party_id: order_id, + maybe_model_type: shared.maybe_model_type, + maybe_prompt_token: shared.maybe_prompt_token, + maybe_wallet_ledger_entry_token, + maybe_creator_user_token: Some(shared.user_token), + maybe_avt_token: shared.maybe_avt_token, + creator_ip_address: shared.ip_address, + creator_set_visibility: Visibility::Public, + maybe_debug_log_event_token: shared.maybe_debug_log_event_token, + mysql_executor: &mut **shared.transaction, + maybe_inference_args: None, + phantom: Default::default(), + } + ).await; + + match db_result { + Ok(token) => all_job_tokens.push(token), + Err(err) => { + warn!("Error inserting seedance2pro inference job (image, order_id={}): {:?}", order_id, err); + if i == 0 { return Err(CommonWebError::from_error(err)); } + } + } + } + + let primary_job_token = all_job_tokens.first().cloned().ok_or_else(|| { + error!("No inference job token was created"); + CommonWebError::server_error_with_message("No inference job token was created") + })?; + + Ok(InsertSeedance2proJobsResult { + primary_job_token, + all_job_tokens, + }) +} diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/insert_db_job/mod.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/insert_db_job/mod.rs new file mode 100644 index 0000000000..751cbf87e7 --- /dev/null +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/insert_db_job/mod.rs @@ -0,0 +1,3 @@ +pub mod insert_fal_job; +pub mod insert_seedance2pro_jobs; +pub mod shared_job_args; diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/insert_db_job/shared_job_args.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/insert_db_job/shared_job_args.rs new file mode 100644 index 0000000000..cc3dbde042 --- /dev/null +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/insert_db_job/shared_job_args.rs @@ -0,0 +1,20 @@ +use enums::common::generation::common_model_type::CommonModelType; +use tokens::tokens::anonymous_visitor_tracking::AnonymousVisitorTrackingToken; +use tokens::tokens::generic_inference_jobs::InferenceJobToken; +use tokens::tokens::non_unique::debug_logs_event_token::DebugLogEventToken; +use tokens::tokens::prompts::PromptToken; +use tokens::tokens::users::UserToken; + +/// Fields every image-side `insert_*_job` helper needs. +/// Parallels `omni_gen::generate::video::insert_db_job::shared_job_args::SharedJobArgs`. +pub struct SharedJobArgs<'a, 'tx> { + pub apriori_job_token: &'a InferenceJobToken, + pub idempotency_token: &'a str, + pub user_token: &'a UserToken, + pub maybe_avt_token: Option<&'a AnonymousVisitorTrackingToken>, + pub maybe_prompt_token: Option<&'a PromptToken>, + pub maybe_debug_log_event_token: Option<&'a DebugLogEventToken>, + pub maybe_model_type: Option, + pub ip_address: &'a str, + pub transaction: &'a mut sqlx::Transaction<'tx, sqlx::MySql>, +} diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/mod.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/mod.rs index 16c7d5276f..54e053a7f5 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/mod.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/mod.rs @@ -1,4 +1,5 @@ pub mod omni_gen_image_generate_handler; pub(crate) mod hydrate_to_router_request; +pub(crate) mod insert_db_job; pub(crate) mod pipeline_result; pub mod pipeline_v2; diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/omni_gen_image_generate_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/omni_gen_image_generate_handler.rs index 3dc08c688b..543db74ea1 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/omni_gen_image_generate_handler.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/omni_gen_image_generate_handler.rs @@ -14,24 +14,23 @@ use enums::by_table::prompts::prompt_type::PromptType; use enums::common::generation::common_generation_mode::CommonGenerationMode; use enums::common::generation::common_model_type::CommonModelType; use enums::common::generation_provider::GenerationProvider; -use enums::common::visibility::Visibility; use http_server_common::request::get_request_ip::get_request_ip; use mysql_queries::queries::debug_logs::insert_debug_log::{insert_debug_log, InsertDebugLogArgs}; -use mysql_queries::queries::generic_inference::api_providers::fal::insert_generic_inference_job_for_fal_queue::FalCategory; -use mysql_queries::queries::generic_inference::api_providers::fal::insert_generic_inference_job_for_fal_queue_with_apriori_job_token::{ - insert_generic_inference_job_for_fal_queue_with_apriori_job_token, - InsertGenericInferenceForFalWithAprioriJobTokenArgs, -}; +use mysql_queries::queries::generic_inference::api_providers::seedance2pro::insert_generic_inference_job_for_seedance2pro_queue_with_apriori_job_token::KinoviVersion; use mysql_queries::queries::idepotency_tokens::insert_idempotency_token::insert_idempotency_token; use mysql_queries::queries::prompt_context_items::insert_batch_prompt_context_items::{ insert_batch_prompt_context_items, InsertBatchArgs, PromptContextItem, }; use mysql_queries::queries::prompts::insert_prompt::{insert_prompt, InsertPromptArgs}; +use tokens::tokens::generic_inference_jobs::InferenceJobToken; use tokens::tokens::non_unique::debug_logs_event_token::DebugLogEventToken; use crate::http_server::common_responses::common_web_error::CommonWebError; use crate::http_server::endpoints::generate::common::payments_error_test::payments_error_test; use crate::http_server::endpoints::omni_gen::generate::image::hydrate_to_router_request::hydrate_to_router_request; +use crate::http_server::endpoints::omni_gen::generate::image::insert_db_job::insert_fal_job::{insert_fal_job, InsertFalJobArgs}; +use crate::http_server::endpoints::omni_gen::generate::image::insert_db_job::insert_seedance2pro_jobs::{insert_seedance2pro_jobs, InsertSeedance2proJobsArgs}; +use crate::http_server::endpoints::omni_gen::generate::image::insert_db_job::shared_job_args::SharedJobArgs; use crate::http_server::endpoints::omni_gen::generate::image::pipeline_v2::run_pipeline_v2::{run_pipeline_v2, RunPipelineV2Args}; use crate::http_server::session::lookup::user_session_feature_flags::UserSessionFeatureFlags; use crate::http_server::validations::validate_idempotency_token_format::validate_idempotency_token_format; @@ -168,20 +167,6 @@ pub async fn omni_gen_image_generate_handler( } } - let external_job_id = match &pipeline_result.response { - GenerateImageResponse::Artcraft(p) => { - p.inference_job_token.as_str().to_string() - } - GenerateImageResponse::Fal(p) => { - p.request_id.clone().unwrap_or_default() - } - GenerateImageResponse::Seedance2Pro(p) => { - // Kinovi/Midjourney uses Seedance2Pro's task_id as the external job - // identifier (the order_id is also available on the payload). - p.task_id.clone() - } - }; - // ==================== WRITE RESULT ==================== // let ip_address = get_request_ip(&http_request); @@ -256,34 +241,63 @@ pub async fn omni_gen_image_generate_handler( } // -- Inference job -- + // + // Each provider has its own queue / worker. Branch by response variant + // so the row lands on the correct queue (the Seedance2Pro/Kinovi worker + // for Midjourney; the Fal worker for everything else). - let db_result = insert_generic_inference_job_for_fal_queue_with_apriori_job_token( - InsertGenericInferenceForFalWithAprioriJobTokenArgs { - apriori_job_token: &pipeline_result.apriori_job_token, - uuid_idempotency_token: &idempotency_token, - maybe_external_third_party_id: &external_job_id, - fal_category: FalCategory::ImageGeneration, - maybe_model_type: request.model.map(|v| v.to_common_model_type()), - maybe_inference_args: None, - maybe_prompt_token: prompt_token.as_ref(), - maybe_creator_user_token: Some(user_token), - maybe_avt_token: maybe_avt_token.as_ref(), - creator_ip_address: &ip_address, - creator_set_visibility: Visibility::Public, - maybe_debug_log_event_token: Some(&debug_log_event_token), - mysql_executor: &mut *transaction, - starting_job_status_override: None, - maybe_frontend_failure_category: None, - maybe_failure_reason: None, - phantom: Default::default(), + let job_token: InferenceJobToken = match &pipeline_result.response { + GenerateImageResponse::Seedance2Pro(payload) => { + info!("Inserting seedance2pro image job(s) with token: {:?}", pipeline_result.apriori_job_token); + + // The image-side omni pipeline always dispatches Midjourney via the + // Volcengine Kinovi account today. If we ever route to BytePlus / + // BytePlus Ultra here, mirror the video-side `kinovi_account` knob. + let kinovi_version = KinoviVersion::Volcengine; + + let result = insert_seedance2pro_jobs(InsertSeedance2proJobsArgs { + primary_order_id: &payload.order_id, + maybe_additional_order_ids: payload.maybe_order_ids.as_deref(), + maybe_wallet_ledger_entry_token: pipeline_result.maybe_wallet_ledger_entry_token.as_ref(), + kinovi_version, + shared: SharedJobArgs { + apriori_job_token: &pipeline_result.apriori_job_token, + idempotency_token: &idempotency_token, + user_token, + maybe_avt_token: maybe_avt_token.as_ref(), + maybe_model_type: request.model.map(|v| v.to_common_model_type()), + maybe_prompt_token: prompt_token.as_ref(), + maybe_debug_log_event_token: Some(&debug_log_event_token), + ip_address: &ip_address, + transaction: &mut transaction, + }, + }).await?; + result.primary_job_token } - ).await; - - let job_token = match db_result { - Ok(token) => token, - Err(err) => { - warn!("Error inserting inference job: {:?}", err); - return Err(CommonWebError::from_error(err)); + GenerateImageResponse::Fal(payload) => { + info!("Inserting fal image job with token: {:?}", pipeline_result.apriori_job_token); + let external_job_id = payload.request_id.clone().unwrap_or_default(); + insert_fal_job(InsertFalJobArgs { + external_job_id: &external_job_id, + shared: SharedJobArgs { + apriori_job_token: &pipeline_result.apriori_job_token, + idempotency_token: &idempotency_token, + user_token, + maybe_avt_token: maybe_avt_token.as_ref(), + maybe_model_type: request.model.map(|v| v.to_common_model_type()), + maybe_prompt_token: prompt_token.as_ref(), + maybe_debug_log_event_token: Some(&debug_log_event_token), + ip_address: &ip_address, + transaction: &mut transaction, + }, + }).await? + } + GenerateImageResponse::Artcraft(payload) => { + // The omni image pipeline never dispatches via the Artcraft provider + // itself today (everything routes to Fal or Kinovi), but the response + // variant exists so we cover it defensively — Artcraft jobs come back + // already inserted server-side, so just propagate the token. + payload.inference_job_token.clone() } }; diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/pipeline_result.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/pipeline_result.rs index 80a094b166..4d4bd9120d 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/pipeline_result.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/pipeline_result.rs @@ -1,8 +1,15 @@ use artcraft_router::generate::generate_image::generate_image_response::GenerateImageResponse; use tokens::tokens::generic_inference_jobs::InferenceJobToken; +use tokens::tokens::wallet_ledger_entries::WalletLedgerEntryToken; /// Both image pipelines produce this, then the shared handler suffix writes DB rows. pub struct ImagePipelineResult { pub apriori_job_token: InferenceJobToken, pub response: GenerateImageResponse, + + /// The wallet ledger entry created when the user was charged for this + /// generation, if any (None when the generation is free). Used as a + /// foreign key on the inference job row so the worker can refund on + /// failure. + pub maybe_wallet_ledger_entry_token: Option, } diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/pipeline_v2/run_pipeline_v2.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/pipeline_v2/run_pipeline_v2.rs index 91ab9145c4..22552a640a 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/pipeline_v2/run_pipeline_v2.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/pipeline_v2/run_pipeline_v2.rs @@ -53,20 +53,24 @@ pub async fn run_pipeline_v2( let apriori_job_token = InferenceJobToken::generate(); - if cost > 0 { - attempt_wallet_deduction_else_common_web_error( + let maybe_wallet_ledger_entry_token = if cost > 0 { + let deduction = attempt_wallet_deduction_else_common_web_error( user_token, Some(apriori_job_token.as_str()), cost, mysql_connection, ).await?; - } + Some(deduction.ledger_entry_token) + } else { + None + }; let response = finalize_and_generate(draft_or_request, server_state, resolved_media).await?; Ok(ImagePipelineResult { apriori_job_token, response, + maybe_wallet_ledger_entry_token, }) }