add correct job type

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<InferenceJobToken, CommonWebError> {
+  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))
+    }
+  }
+}

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<InferenceJobToken>,
+}
+
+/// 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<InsertSeedance2proJobsResult, CommonWebError> {
+  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<InferenceJobToken> = 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,
+  })
+}

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;

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<CommonModelType>,
+  pub ip_address: &'a str,
+  pub transaction: &'a mut sqlx::Transaction<'tx, sqlx::MySql>,
+}

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;

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()
     }
   };
 

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<WalletLedgerEntryToken>,
 }

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,
   })
 }